1package zerolog
2
3import (
4	"bytes"
5	"io"
6	"path"
7	"runtime"
8	"strconv"
9	"strings"
10	"sync"
11)
12
13// LevelWriter defines as interface a writer may implement in order
14// to receive level information with payload.
15type LevelWriter interface {
16	io.Writer
17	WriteLevel(level Level, p []byte) (n int, err error)
18}
19
20type levelWriterAdapter struct {
21	io.Writer
22}
23
24func (lw levelWriterAdapter) WriteLevel(l Level, p []byte) (n int, err error) {
25	return lw.Write(p)
26}
27
28type syncWriter struct {
29	mu sync.Mutex
30	lw LevelWriter
31}
32
33// SyncWriter wraps w so that each call to Write is synchronized with a mutex.
34// This syncer can be used to wrap the call to writer's Write method if it is
35// not thread safe. Note that you do not need this wrapper for os.File Write
36// operations on POSIX and Windows systems as they are already thread-safe.
37func SyncWriter(w io.Writer) io.Writer {
38	if lw, ok := w.(LevelWriter); ok {
39		return &syncWriter{lw: lw}
40	}
41	return &syncWriter{lw: levelWriterAdapter{w}}
42}
43
44// Write implements the io.Writer interface.
45func (s *syncWriter) Write(p []byte) (n int, err error) {
46	s.mu.Lock()
47	defer s.mu.Unlock()
48	return s.lw.Write(p)
49}
50
51// WriteLevel implements the LevelWriter interface.
52func (s *syncWriter) WriteLevel(l Level, p []byte) (n int, err error) {
53	s.mu.Lock()
54	defer s.mu.Unlock()
55	return s.lw.WriteLevel(l, p)
56}
57
58type multiLevelWriter struct {
59	writers []LevelWriter
60}
61
62func (t multiLevelWriter) Write(p []byte) (n int, err error) {
63	for _, w := range t.writers {
64		if _n, _err := w.Write(p); err == nil {
65			n = _n
66			if _err != nil {
67				err = _err
68			} else if _n != len(p) {
69				err = io.ErrShortWrite
70			}
71		}
72	}
73	return n, err
74}
75
76func (t multiLevelWriter) WriteLevel(l Level, p []byte) (n int, err error) {
77	for _, w := range t.writers {
78		if _n, _err := w.WriteLevel(l, p); err == nil {
79			n = _n
80			if _err != nil {
81				err = _err
82			} else if _n != len(p) {
83				err = io.ErrShortWrite
84			}
85		}
86	}
87	return n, err
88}
89
90// MultiLevelWriter creates a writer that duplicates its writes to all the
91// provided writers, similar to the Unix tee(1) command. If some writers
92// implement LevelWriter, their WriteLevel method will be used instead of Write.
93func MultiLevelWriter(writers ...io.Writer) LevelWriter {
94	lwriters := make([]LevelWriter, 0, len(writers))
95	for _, w := range writers {
96		if lw, ok := w.(LevelWriter); ok {
97			lwriters = append(lwriters, lw)
98		} else {
99			lwriters = append(lwriters, levelWriterAdapter{w})
100		}
101	}
102	return multiLevelWriter{lwriters}
103}
104
105// TestingLog is the logging interface of testing.TB.
106type TestingLog interface {
107	Log(args ...interface{})
108	Logf(format string, args ...interface{})
109	Helper()
110}
111
112// TestWriter is a writer that writes to testing.TB.
113type TestWriter struct {
114	T TestingLog
115
116	// Frame skips caller frames to capture the original file and line numbers.
117	Frame int
118}
119
120// NewTestWriter creates a writer that logs to the testing.TB.
121func NewTestWriter(t TestingLog) TestWriter {
122	return TestWriter{T: t}
123}
124
125// Write to testing.TB.
126func (t TestWriter) Write(p []byte) (n int, err error) {
127	t.T.Helper()
128
129	n = len(p)
130
131	// Strip trailing newline because t.Log always adds one.
132	p = bytes.TrimRight(p, "\n")
133
134	// Try to correct the log file and line number to the caller.
135	if t.Frame > 0 {
136		_, origFile, origLine, _ := runtime.Caller(1)
137		_, frameFile, frameLine, ok := runtime.Caller(1 + t.Frame)
138		if ok {
139			erase := strings.Repeat("\b", len(path.Base(origFile))+len(strconv.Itoa(origLine))+3)
140			t.T.Logf("%s%s:%d: %s", erase, path.Base(frameFile), frameLine, p)
141			return n, err
142		}
143	}
144	t.T.Log(string(p))
145
146	return n, err
147}
148
149// ConsoleTestWriter creates an option that correctly sets the file frame depth for testing.TB log.
150func ConsoleTestWriter(t TestingLog) func(w *ConsoleWriter) {
151	return func(w *ConsoleWriter) {
152		w.Out = TestWriter{T: t, Frame: 6}
153	}
154}
155