From 2a4dfabd8db65ab168c6f0f734c88b4db374a15b Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Thu, 12 Sep 2024 16:48:25 +0200 Subject: [PATCH 1/7] Limit maximum number of connections --- server.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/server.go b/server.go index 7c23a31..4cccbe0 100644 --- a/server.go +++ b/server.go @@ -32,6 +32,7 @@ type Server struct { LMTP bool Domain string + MaxConnections int MaxRecipients int MaxMessageBytes int64 MaxLineLength int @@ -151,6 +152,12 @@ func (s *Server) handleConn(c *Conn) error { } } + // limit connections + if s.MaxConnections > 0 && len(s.conns) > s.MaxConnections { + c.Reject() + return nil + } + c.greet() for { From 22834b4e0077555a8fd340d99f295e0ca0678295 Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Wed, 18 Sep 2024 10:48:39 +0200 Subject: [PATCH 2/7] Fix missing s.locker.lock() and duplicate c.Close() --- server.go | 35 +++++++++++++++++++++-------------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/server.go b/server.go index 4cccbe0..801a3b7 100644 --- a/server.go +++ b/server.go @@ -128,17 +128,7 @@ func (s *Server) Serve(l net.Listener) error { } func (s *Server) handleConn(c *Conn) error { - s.locker.Lock() - s.conns[c] = struct{}{} - s.locker.Unlock() - - defer func() { - c.Close() - - s.locker.Lock() - delete(s.conns, c) - s.locker.Unlock() - }() + defer c.Close() if tlsConn, ok := c.conn.(*tls.Conn); ok { if d := s.ReadTimeout; d != 0 { @@ -153,11 +143,28 @@ func (s *Server) handleConn(c *Conn) error { } // limit connections - if s.MaxConnections > 0 && len(s.conns) > s.MaxConnections { - c.Reject() - return nil + if s.MaxConnections > 0 { + s.locker.Lock() + nConns := len(s.conns) + s.locker.Unlock() + + if nConns > s.MaxConnections { + c.writeResponse(421, EnhancedCode{4, 4, 5}, "Too busy. Try again later.") + return nil + } } + // register connection + s.locker.Lock() + s.conns[c] = struct{}{} + s.locker.Unlock() + + defer func() { + s.locker.Lock() + delete(s.conns, c) + s.locker.Unlock() + }() + c.greet() for { From 8ff94a73745ce1d44c77d32e0b98b4a04b9ab514 Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Wed, 18 Sep 2024 10:59:44 +0200 Subject: [PATCH 3/7] Fix race condition --- server.go | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/server.go b/server.go index 801a3b7..64d4e88 100644 --- a/server.go +++ b/server.go @@ -142,23 +142,23 @@ func (s *Server) handleConn(c *Conn) error { } } - // limit connections - if s.MaxConnections > 0 { - s.locker.Lock() - nConns := len(s.conns) - s.locker.Unlock() - - if nConns > s.MaxConnections { - c.writeResponse(421, EnhancedCode{4, 4, 5}, "Too busy. Try again later.") - return nil - } - } - // register connection + maxConnsExceeded := false s.locker.Lock() - s.conns[c] = struct{}{} + if s.MaxConnections > 0 && len(s.conns) >= s.MaxConnections { + maxConnsExceeded = true + } else { + s.conns[c] = struct{}{} + } s.locker.Unlock() + // limit connections + if maxConnsExceeded { + c.writeResponse(421, EnhancedCode{4, 4, 5}, "Too busy. Try again later.") + return nil + } + + // unregister connection defer func() { s.locker.Lock() delete(s.conns, c) From 2ae8ba2b2b81ecd711f3d34ce18a60fb0a5bd10f Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Wed, 18 Sep 2024 11:30:22 +0200 Subject: [PATCH 4/7] Add unit test --- server_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/server_test.go b/server_test.go index 206d336..f5ef381 100644 --- a/server_test.go +++ b/server_test.go @@ -14,6 +14,7 @@ import ( "testing" "github.com/emersion/go-sasl" + "github.com/emersion/go-smtp" ) @@ -1514,3 +1515,49 @@ func TestServerDSNwithSMTPUTF8(t *testing.T) { t.Fatal("Invalid ORCPT address:", val) } } + +func TestServer_MaxConnections(t *testing.T) { + cases := []struct { + name string + maxConnections int + expected string + }{ + // 0 = unlimited; all connections should be accepted + {name: "MaxConnections set to 0", maxConnections: 0, expected: "220 localhost ESMTP Service Ready"}, + // 1 = only one connection is allowed; the second connection should be rejected + {name: "MaxConnections set to 1", maxConnections: 1, expected: "421 4.4.5 Too busy. Try again later."}, + // 2 = two connections are allowed; the second connection should be accepted + {name: "MaxConnections set to 2", maxConnections: 2, expected: "220 localhost ESMTP Service Ready"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + // create server with a single allowed connection + _, s, c, scanner1 := testServer(t, func(s *smtp.Server) { + s.MaxConnections = tc.maxConnections + }) + defer s.Close() + + // there is already be one connection registered + // and we can read the greeting from it (see testServerGreeted()) + scanner1.Scan() + if scanner1.Text() != "220 localhost ESMTP Service Ready" { + t.Fatal("Invalid first greeting:", scanner1.Text()) + } + + // now we create a second connection + c2, err := net.Dial("tcp", c.RemoteAddr().String()) + if err != nil { + t.Fatal("Error creating second connection:", err) + } + + // we should get an 421 error greeting now + scanner2 := bufio.NewScanner(c2) + scanner2.Scan() + if scanner2.Text() != tc.expected { + t.Fatal("Invalid second greeting:", scanner2.Text()) + } + }) + } + +} From 94fcc072a1d1c6b1976db855c1338b14e79a9bf7 Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Wed, 18 Sep 2024 11:31:20 +0200 Subject: [PATCH 5/7] revert unnecessary change --- server_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/server_test.go b/server_test.go index f5ef381..e935552 100644 --- a/server_test.go +++ b/server_test.go @@ -14,7 +14,6 @@ import ( "testing" "github.com/emersion/go-sasl" - "github.com/emersion/go-smtp" ) From b1ea4cc967d05ac3b9bbf18b51e7740fa7cdf142 Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Wed, 18 Sep 2024 11:45:13 +0200 Subject: [PATCH 6/7] fix test doc --- server_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server_test.go b/server_test.go index e935552..657a24b 100644 --- a/server_test.go +++ b/server_test.go @@ -1550,7 +1550,7 @@ func TestServer_MaxConnections(t *testing.T) { t.Fatal("Error creating second connection:", err) } - // we should get an 421 error greeting now + // we should get an appropriate greeting now scanner2 := bufio.NewScanner(c2) scanner2.Scan() if scanner2.Text() != tc.expected { From 6c510f06a10328981333cf9a2fe49c195dbc2493 Mon Sep 17 00:00:00 2001 From: Mathias Lieber Date: Wed, 18 Sep 2024 11:56:29 +0200 Subject: [PATCH 7/7] fix test doc --- server_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server_test.go b/server_test.go index 657a24b..f5cd3ca 100644 --- a/server_test.go +++ b/server_test.go @@ -1531,7 +1531,7 @@ func TestServer_MaxConnections(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - // create server with a single allowed connection + // create server with limited allowed connections _, s, c, scanner1 := testServer(t, func(s *smtp.Server) { s.MaxConnections = tc.maxConnections })