@@ -2,129 +2,90 @@ package agent
22
33import (
44 "bytes"
5- "context"
6- "fmt"
75 "io"
86 "sync"
7+ "sync/atomic"
98
10- "github.com/fnproject/fn/api/common"
119 "github.com/fnproject/fn/api/models"
1210 "github.com/sirupsen/logrus"
1311)
1412
1513var (
1614 bufPool = & sync.Pool {New : func () interface {} { return new (bytes.Buffer ) }}
17- logPool = & sync.Pool {New : func () interface {} { return new (bytes.Buffer ) }}
1815)
1916
20- // setupLogger returns a ReadWriteCloser that may have:
21- // * [always] writes bytes to a size limited buffer, that can be read from using io.Reader
2217// * [always] writes bytes per line to stderr as DEBUG
23- //
24- // To prevent write failures from failing the call or any other writes,
25- // multiWriteCloser ignores errors. Close will flush the line writers
26- // appropriately. The returned io.ReadWriteCloser is not safe for use after
27- // calling Close.
28- func setupLogger (ctx context.Context , maxSize uint64 , debug bool , c * models.Call ) io.ReadWriteCloser {
18+ func setupLogger (c * models.Call ) io.WriteCloser {
2919 lbuf := bufPool .Get ().(* bytes.Buffer )
30- dbuf := logPool .Get ().(* bytes.Buffer )
3120
32- close := func () error {
21+ close := func () {
3322 // TODO we may want to toss out buffers that grow to grotesque size but meh they will prob get GC'd
3423 lbuf .Reset ()
35- dbuf .Reset ()
3624 bufPool .Put (lbuf )
37- logPool .Put (dbuf )
38- return nil
3925 }
4026
41- // we don't need to log per line to db, but we do need to limit it
42- limitw := & nopCloser {newLimitWriter (int (maxSize ), dbuf )}
43-
44- // order matters, in that closer should be last and limit should be next to last
45- mw := make (multiWriteCloser , 0 , 3 )
46-
47- if debug {
48- // accumulate all line writers, wrap in same line writer (to re-use buffer)
49- stderrLogger := common .Logger (ctx ).WithFields (logrus.Fields {"user_log" : true , "app_id" : c .AppID , "fn_id" : c .FnID , "image" : c .Image , "call_id" : c .ID })
50- loggo := & nopCloser {& logWriter {stderrLogger }}
51-
52- // we don't need to limit the log writer(s), but we do need it to dispense lines
53- linew := newLineWriterWithBuffer (lbuf , loggo )
54- mw = append (mw , linew )
27+ stderrLogger := logrus .WithFields (logrus.Fields {"user_log" : true , "app_id" : c .AppID , "fn_id" : c .FnID , "image" : c .Image , "call_id" : c .ID })
28+ loggo := newLogWriter (stderrLogger )
29+ linew := newLineWriterWithBuffer (lbuf , loggo )
30+ linew = & fCloser {
31+ close : func () error {
32+ err := linew .Close ()
33+ close ()
34+ return err
35+ },
5536 }
56-
57- mw = append (mw , limitw , & fCloser {close })
58- return & rwc {mw , dbuf }
37+ return linew
5938}
6039
61- // implements io.ReadWriteCloser, fmt.Stringer and Bytes()
62- // TODO WriteString and ReadFrom would be handy to implement,
63- // ReadFrom is a little involved.
64- type rwc struct {
65- io.WriteCloser
66-
67- // buffer is not embedded since it would bypass calls to WriteCloser.Write
68- // in cases such as WriteString and ReadFrom
69- b * bytes.Buffer
70- }
71-
72- func (r * rwc ) Read (b []byte ) (int , error ) { return r .b .Read (b ) }
73- func (r * rwc ) String () string { return r .b .String () }
74- func (r * rwc ) Bytes () []byte { return r .b .Bytes () }
75-
76- // implements passthrough Write & closure call in Close
40+ // implements passthrough WriteCloser with overwritable Close
7741type fCloser struct {
42+ io.Writer
7843 close func () error
7944}
8045
81- func (f * fCloser ) Write (b []byte ) (int , error ) { return len (b ), nil }
82- func (f * fCloser ) Close () error { return f .close () }
46+ func (f * fCloser ) Close () error { return f .close () }
8347
8448type nopCloser struct {
8549 io.Writer
8650}
8751
8852func (n * nopCloser ) Close () error { return nil }
8953
90- // multiWriteCloser ignores all errors from inner writers. you say, oh, this is a bad idea?
91- // yes, well, we were going to silence them all individually anyway, so let's not be shy about it.
92- // the main thing we need to ensure is that every close is called, even if another errors.
93- // XXX(reed): maybe we should log it (for syslog, it may help debug, maybe we just log that one)
94- type multiWriteCloser []io.WriteCloser
95-
96- func (m multiWriteCloser ) Write (b []byte ) (n int , err error ) {
97- for _ , mw := range m {
98- mw .Write (b )
99- }
100- return len (b ), nil
101- }
102-
103- func (m multiWriteCloser ) Close () (err error ) {
104- for _ , mw := range m {
105- mw .Close ()
106- }
107- return nil
108- }
109-
11054// logWriter will log (to real stderr) every call to Write as a line. it should
11155// be wrapped with a lineWriter so that the output makes sense.
11256type logWriter struct {
57+ // level string // XXX(reed):
11358 logrus.FieldLogger
59+ closed uint32
60+ }
61+
62+ func newLogWriter (logger logrus.FieldLogger ) io.WriteCloser {
63+ return & logWriter {FieldLogger : logger }
11464}
11565
11666func (l * logWriter ) Write (b []byte ) (int , error ) {
67+ if atomic .LoadUint32 (& l .closed ) == 1 {
68+ // we don't want to return 0/error or the container will get shut down
69+ return len (b ), nil
70+ }
11771 l .Debug (string (b ))
11872 return len (b ), nil
11973}
12074
75+ func (l * logWriter ) Close () error {
76+ atomic .StoreUint32 (& l .closed , 1 )
77+ return nil
78+ }
79+
12180// lineWriter buffers all calls to Write and will call Write
12281// on the underlying writer once per new line. Close must
12382// be called to ensure that the buffer is flushed, and a newline
12483// will be appended in Close if none is present.
84+ // TODO(reed): is line writer is vulnerable to attack?
12585type lineWriter struct {
126- b * bytes.Buffer
127- w io.WriteCloser
86+ b * bytes.Buffer
87+ w io.WriteCloser
88+ closed uint32
12889}
12990
13091func newLineWriter (w io.WriteCloser ) io.WriteCloser {
@@ -136,6 +97,10 @@ func newLineWriterWithBuffer(b *bytes.Buffer, w io.WriteCloser) io.WriteCloser {
13697}
13798
13899func (li * lineWriter ) Write (ogb []byte ) (int , error ) {
100+ if atomic .LoadUint32 (& li .closed ) == 1 {
101+ // we don't want to return 0/error or the container will shut down
102+ return len (ogb ), nil
103+ }
139104 li .b .Write (ogb ) // bytes.Buffer is guaranteed, read it!
140105
141106 for {
@@ -159,6 +124,8 @@ func (li *lineWriter) Write(ogb []byte) (int, error) {
159124}
160125
161126func (li * lineWriter ) Close () error {
127+ atomic .StoreUint32 (& li .closed , 1 )
128+
162129 defer li .w .Close () // MUST close this (after writing last line)
163130
164131 // flush the remaining bytes in the buffer to underlying writer, adding a
@@ -174,41 +141,3 @@ func (li *lineWriter) Close() error {
174141 _ , err := li .w .Write (b )
175142 return err
176143}
177-
178- // io.Writer that allows limiting bytes written to w
179- // TODO change to use clamp writer, this is dupe code
180- type limitDiscardWriter struct {
181- n , max int
182- io.Writer
183- }
184-
185- func newLimitWriter (max int , w io.Writer ) io.Writer {
186- return & limitDiscardWriter {max : max , Writer : w }
187- }
188-
189- func (l * limitDiscardWriter ) Write (b []byte ) (int , error ) {
190- inpLen := len (b )
191- if l .n >= l .max {
192- return inpLen , nil
193- }
194-
195- if l .n + inpLen >= l .max {
196- // cut off to prevent gigantic line attack
197- b = b [:l .max - l .n ]
198- }
199-
200- n , err := l .Writer .Write (b )
201- l .n += n
202-
203- if l .n >= l .max {
204- // write in truncation message to log once
205- l .Writer .Write ([]byte (fmt .Sprintf ("\n -----max log size %d bytes exceeded, truncating log-----\n " , l .max )))
206- } else if n != len (b ) {
207- // Is this truly a partial write? We'll be honest if that's the case.
208- return n , err
209- }
210-
211- // yes, we lie... this is to prevent callers to blow up, we always pretend
212- // that we were able to write the entire buffer.
213- return inpLen , err
214- }
0 commit comments