From f804d2f90f9eeb202f5de693d400cc00b7a97ee5 Mon Sep 17 00:00:00 2001
From: Nathan Smith <12156185+nsmith5@users.noreply.github.com>
Date: Sun, 13 Sep 2020 11:23:32 -0700
Subject: [PATCH] server: Improved error handling (#116)

* Factor error handling into seperate function

* Rename error handler and change status for too many errors

* Make error threshold a global const

* Switches protocolError status codes to 500 5.5.1

* Patches too many invalid commands test

Now matches the 500 code thrown by protocolError
---
 conn.go   | 45 +++++++++++++++++++++++++++------------------
 server.go |  3 +--
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/conn.go b/conn.go
index 0ca88a1..961d863 100644
--- a/conn.go
+++ b/conn.go
@@ -17,6 +17,9 @@ import (
 	"time"
 )
 
+// Number of errors we'll tolerate per connection before closing. Defaults to 3.
+const errThreshold = 3
+
 type ConnectionState struct {
 	Hostname   string
 	LocalAddr  net.Addr
@@ -25,11 +28,14 @@ type ConnectionState struct {
 }
 
 type Conn struct {
-	conn       net.Conn
-	text       *textproto.Conn
-	server     *Server
-	helo       string
-	nbrErrors  int
+	conn   net.Conn
+	text   *textproto.Conn
+	server *Server
+	helo   string
+
+	// Number of errors witnessed on this connection
+	errCount int
+
 	session    Session
 	locker     sync.Mutex
 	binarymime bool
@@ -82,16 +88,6 @@ func (c *Conn) init() {
 	c.text = textproto.NewConn(rwc)
 }
 
-func (c *Conn) unrecognizedCommand(cmd string) {
-	c.WriteResponse(500, EnhancedCode{5, 5, 2}, fmt.Sprintf("Syntax error, %v command unrecognized", cmd))
-
-	c.nbrErrors++
-	if c.nbrErrors > 3 {
-		c.WriteResponse(500, EnhancedCode{5, 5, 2}, "Too many unrecognized commands")
-		c.Close()
-	}
-}
-
 // Commands are dispatched to the appropriate handler functions.
 func (c *Conn) handle(cmd string, arg string) {
 	// If panic happens during command handling - send 421 response
@@ -107,7 +103,7 @@ func (c *Conn) handle(cmd string, arg string) {
 	}()
 
 	if cmd == "" {
-		c.WriteResponse(500, EnhancedCode{5, 5, 2}, "Speak up")
+		c.protocolError(500, EnhancedCode{5, 5, 2}, "Speak up")
 		return
 	}
 
@@ -148,14 +144,15 @@ func (c *Conn) handle(cmd string, arg string) {
 		c.Close()
 	case "AUTH":
 		if c.server.AuthDisabled {
-			c.unrecognizedCommand(cmd)
+			c.protocolError(500, EnhancedCode{5, 5, 2}, "Syntax error, AUTH command unrecognized")
 		} else {
 			c.handleAuth(arg)
 		}
 	case "STARTTLS":
 		c.handleStartTLS()
 	default:
-		c.unrecognizedCommand(cmd)
+		msg := fmt.Sprintf("Syntax errors, %v command unrecognized", cmd)
+		c.protocolError(500, EnhancedCode{5, 5, 2}, msg)
 	}
 }
 
@@ -222,6 +219,18 @@ func (c *Conn) authAllowed() bool {
 	return !c.server.AuthDisabled && (isTLS || c.server.AllowInsecureAuth)
 }
 
+// protocolError writes errors responses and closes the connection once too many
+// have occurred.
+func (c *Conn) protocolError(code int, ec EnhancedCode, msg string) {
+	c.WriteResponse(code, ec, msg)
+
+	c.errCount++
+	if c.errCount > errThreshold {
+		c.WriteResponse(500, EnhancedCode{5, 5, 1}, "Too many errors. Quiting now")
+		c.Close()
+	}
+}
+
 // GREET state -> waiting for HELO
 func (c *Conn) handleGreet(enhanced bool, arg string) {
 	if !enhanced {
diff --git a/server.go b/server.go
index 52506b5..a7c37c9 100644
--- a/server.go
+++ b/server.go
@@ -147,8 +147,7 @@ func (s *Server) handleConn(c *Conn) error {
 		if err == nil {
 			cmd, arg, err := parseCmd(line)
 			if err != nil {
-				c.nbrErrors++
-				c.WriteResponse(501, EnhancedCode{5, 5, 2}, "Bad command")
+				c.protocolError(501, EnhancedCode{5, 5, 2}, "Bad command")
 				continue
 			}