1package writer
 2
 3import (
 4	"bytes"
 5	"io"
 6	"sync"
 7)
 8
 9type WriterInterface interface {
10	io.Writer
11
12	Truncate()
13	DumpOut()
14	DumpOutWithHeader(header string)
15	Bytes() []byte
16}
17
18type Writer struct {
19	buffer     *bytes.Buffer
20	outWriter  io.Writer
21	lock       *sync.Mutex
22	stream     bool
23	redirector io.Writer
24}
25
26func New(outWriter io.Writer) *Writer {
27	return &Writer{
28		buffer:    &bytes.Buffer{},
29		lock:      &sync.Mutex{},
30		outWriter: outWriter,
31		stream:    true,
32	}
33}
34
35func (w *Writer) AndRedirectTo(writer io.Writer) {
36	w.redirector = writer
37}
38
39func (w *Writer) SetStream(stream bool) {
40	w.lock.Lock()
41	defer w.lock.Unlock()
42	w.stream = stream
43}
44
45func (w *Writer) Write(b []byte) (n int, err error) {
46	w.lock.Lock()
47	defer w.lock.Unlock()
48
49	n, err = w.buffer.Write(b)
50	if w.redirector != nil {
51		w.redirector.Write(b)
52	}
53	if w.stream {
54		return w.outWriter.Write(b)
55	}
56	return n, err
57}
58
59func (w *Writer) Truncate() {
60	w.lock.Lock()
61	defer w.lock.Unlock()
62	w.buffer.Reset()
63}
64
65func (w *Writer) DumpOut() {
66	w.lock.Lock()
67	defer w.lock.Unlock()
68	if !w.stream {
69		w.buffer.WriteTo(w.outWriter)
70	}
71}
72
73func (w *Writer) Bytes() []byte {
74	w.lock.Lock()
75	defer w.lock.Unlock()
76	b := w.buffer.Bytes()
77	copied := make([]byte, len(b))
78	copy(copied, b)
79	return copied
80}
81
82func (w *Writer) DumpOutWithHeader(header string) {
83	w.lock.Lock()
84	defer w.lock.Unlock()
85	if !w.stream && w.buffer.Len() > 0 {
86		w.outWriter.Write([]byte(header))
87		w.buffer.WriteTo(w.outWriter)
88	}
89}