From 9dbffa88de3829fe715dbcea012378d5e84b779b Mon Sep 17 00:00:00 2001 From: James Hillyerd Date: Wed, 24 Aug 2022 21:42:53 -0700 Subject: [PATCH] Refactor SMTP handler_test log collection (#289) * smtp: allow logger to be passed into startSession * smtp: remove logbuf and teardown from handler_test * smtp: handler_test log output to testing.T --- pkg/server/smtp/handler.go | 5 +- pkg/server/smtp/handler_test.go | 102 +++++++++++--------------------- pkg/server/smtp/listener.go | 2 +- 3 files changed, 36 insertions(+), 73 deletions(-) diff --git a/pkg/server/smtp/handler.go b/pkg/server/smtp/handler.go index f9e9a55..cc1c20a 100644 --- a/pkg/server/smtp/handler.go +++ b/pkg/server/smtp/handler.go @@ -15,7 +15,6 @@ import ( "github.com/inbucket/inbucket/pkg/policy" "github.com/rs/zerolog" - "github.com/rs/zerolog/log" ) // State tracks the current mode of our SMTP state machine. @@ -144,8 +143,8 @@ func (s *Session) String() string { * 4. If bad cmd, respond error * 5. Goto 2 */ -func (s *Server) startSession(id int, conn net.Conn) { - logger := log.Hook(logHook{}).With(). +func (s *Server) startSession(id int, conn net.Conn, logger zerolog.Logger) { + logger = logger.Hook(logHook{}).With(). Str("module", "smtp"). Str("remote", conn.RemoteAddr().String()). Int("session", id).Logger() diff --git a/pkg/server/smtp/handler_test.go b/pkg/server/smtp/handler_test.go index f07e824..e1f67e1 100644 --- a/pkg/server/smtp/handler_test.go +++ b/pkg/server/smtp/handler_test.go @@ -1,14 +1,11 @@ package smtp import ( - "bytes" "fmt" "io" - "log" "net" "net/textproto" - "os" "testing" "time" @@ -17,6 +14,7 @@ import ( "github.com/inbucket/inbucket/pkg/policy" "github.com/inbucket/inbucket/pkg/storage" "github.com/inbucket/inbucket/pkg/test" + "github.com/rs/zerolog" ) type scriptStep struct { @@ -27,8 +25,8 @@ type scriptStep struct { // Test commands in GREET state func TestGreetState(t *testing.T) { ds := test.NewStore() - server, logbuf, teardown := setupSMTPServer(ds) - defer teardown() + server := setupSMTPServer(ds) + defer server.Drain() // Required to prevent test logging data race. // Test out some mangled HELOs script := []scriptStep{ @@ -77,19 +75,13 @@ func TestGreetState(t *testing.T) { t.Error(err) } - if t.Failed() { - // Wait for handler to finish logging - time.Sleep(2 * time.Second) - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) - } } // Test commands in READY state func TestEmptyEnvelope(t *testing.T) { ds := test.NewStore() - server, logbuf, teardown := setupSMTPServer(ds) - defer teardown() + server := setupSMTPServer(ds) + defer server.Drain() // Test out some empty envelope without blanks script := []scriptStep{ @@ -97,8 +89,6 @@ func TestEmptyEnvelope(t *testing.T) { {"MAIL FROM:<>", 250}, } if err := playSession(t, server, script); err != nil { - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) t.Error(err) } @@ -108,8 +98,6 @@ func TestEmptyEnvelope(t *testing.T) { {"MAIL FROM: <>", 250}, } if err := playSession(t, server, script); err != nil { - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) t.Error(err) } } @@ -117,8 +105,8 @@ func TestEmptyEnvelope(t *testing.T) { // Test AUTH func TestAuth(t *testing.T) { ds := test.NewStore() - server, logbuf, teardown := setupSMTPServer(ds) - defer teardown() + server := setupSMTPServer(ds) + defer server.Drain() // PLAIN AUTH script := []scriptStep{ @@ -149,20 +137,13 @@ func TestAuth(t *testing.T) { if err := playSession(t, server, script); err != nil { t.Error(err) } - - if t.Failed() { - // Wait for handler to finish logging - time.Sleep(2 * time.Second) - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) - } } // Test commands in READY state func TestReadyState(t *testing.T) { ds := test.NewStore() - server, logbuf, teardown := setupSMTPServer(ds) - defer teardown() + server := setupSMTPServer(ds) + defer server.Drain() // Test out some mangled READY commands script := []scriptStep{ @@ -221,20 +202,13 @@ func TestReadyState(t *testing.T) { if err := playSession(t, server, script); err != nil { t.Error(err) } - - if t.Failed() { - // Wait for handler to finish logging - time.Sleep(2 * time.Second) - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) - } } // Test commands in MAIL state func TestMailState(t *testing.T) { mds := test.NewStore() - server, logbuf, teardown := setupSMTPServer(mds) - defer teardown() + server := setupSMTPServer(mds) + defer server.Drain() // Test out some mangled READY commands script := []scriptStep{ @@ -335,23 +309,16 @@ func TestMailState(t *testing.T) { if err := playSession(t, server, script); err != nil { t.Error(err) } - - if t.Failed() { - // Wait for handler to finish logging - time.Sleep(2 * time.Second) - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) - } } // Test commands in DATA state func TestDataState(t *testing.T) { mds := test.NewStore() - server, logbuf, teardown := setupSMTPServer(mds) - defer teardown() + server := setupSMTPServer(mds) + defer server.Drain() var script []scriptStep - pipe := setupSMTPSession(server) + pipe := setupSMTPSession(t, server) c := textproto.NewConn(pipe) if code, _, err := c.ReadCodeLine(220); err != nil { @@ -366,6 +333,7 @@ func TestDataState(t *testing.T) { if err := playScriptAgainst(t, c, script); err != nil { t.Error(err) } + // Send a message body := `To: u1@gmail.com From: john@gmail.com @@ -379,9 +347,11 @@ Hi! if code, _, err := c.ReadCodeLine(250); err != nil { t.Errorf("Expected a 250 greeting, got %v", code) } + _, _ = c.Cmd("QUIT") + _, _, _ = c.ReadCodeLine(221) // Test with no useful headers. - pipe = setupSMTPSession(server) + pipe = setupSMTPSession(t, server) c = textproto.NewConn(pipe) if code, _, err := c.ReadCodeLine(220); err != nil { t.Errorf("Expected a 220 greeting, got %v", code) @@ -395,29 +365,25 @@ Hi! if err := playScriptAgainst(t, c, script); err != nil { t.Error(err) } + // Send a message body = `X-Useless-Header: true -Hi! Can you still deliver this? -` + Hi! Can you still deliver this? + ` dw = c.DotWriter() _, _ = io.WriteString(dw, body) _ = dw.Close() if code, _, err := c.ReadCodeLine(250); err != nil { t.Errorf("Expected a 250 greeting, got %v", code) } - - if t.Failed() { - // Wait for handler to finish logging - time.Sleep(2 * time.Second) - // Dump buffered log data if there was a failure - _, _ = io.Copy(os.Stderr, logbuf) - } + _, _ = c.Cmd("QUIT") + _, _, _ = c.ReadCodeLine(221) } // playSession creates a new session, reads the greeting and then plays the script func playSession(t *testing.T, server *Server, script []scriptStep) error { - pipe := setupSMTPSession(server) + pipe := setupSMTPSession(t, server) c := textproto.NewConn(pipe) if code, _, err := c.ReadCodeLine(220); err != nil { @@ -467,7 +433,7 @@ func (m *mockConn) SetDeadline(t time.Time) error { return nil } func (m *mockConn) SetReadDeadline(t time.Time) error { return nil } func (m *mockConn) SetWriteDeadline(t time.Time) error { return nil } -func setupSMTPServer(ds storage.Store) (s *Server, buf *bytes.Buffer, teardown func()) { +func setupSMTPServer(ds storage.Store) *Server { cfg := &config.Root{ MailboxNaming: config.FullNaming, SMTP: config.SMTP{ @@ -480,26 +446,24 @@ func setupSMTPServer(ds storage.Store) (s *Server, buf *bytes.Buffer, teardown f Timeout: 5, }, } - // Capture log output. - buf = new(bytes.Buffer) - log.SetOutput(buf) + // Create a server, don't start it. - // TODO Remove teardown. - teardown = func() {} addrPolicy := &policy.Addressing{Config: cfg} manager := &message.StoreManager{Store: ds} - s = NewServer(cfg.SMTP, manager, addrPolicy) - return s, buf, teardown + + return NewServer(cfg.SMTP, manager, addrPolicy) } var sessionNum int -func setupSMTPSession(server *Server) net.Conn { - // Pair of pipes to communicate. +func setupSMTPSession(t *testing.T, server *Server) net.Conn { + logger := zerolog.New(zerolog.NewTestWriter(t)) serverConn, clientConn := net.Pipe() + // Start the session. server.wg.Add(1) sessionNum++ - go server.startSession(sessionNum, &mockConn{serverConn}) + go server.startSession(sessionNum, &mockConn{serverConn}, logger) + return clientConn } diff --git a/pkg/server/smtp/listener.go b/pkg/server/smtp/listener.go index 5a281f7..0a38282 100644 --- a/pkg/server/smtp/listener.go +++ b/pkg/server/smtp/listener.go @@ -170,7 +170,7 @@ func (s *Server) serve(ctx context.Context) { tempDelay = 0 expConnectsTotal.Add(1) s.wg.Add(1) - go s.startSession(sessionID, conn) + go s.startSession(sessionID, conn, log.Logger) } } }