refactor(pop3): extract transaction handler helpers (#619)

- validateMsgNum parses and range-checks message-number arguments,
  replacing five copies of the parse/check/respond sequence in the
  LIST, UIDL, DELE, RETR and TOP arms (wire responses unchanged).
- listMessages implements LIST and UIDL, which differ only in the
  per-message value reported.
- sendMessageLines unifies sendMessage/sendMessageTop via a body line
  limit, negative meaning unlimited.

Adds a direct unit test for validateMsgNum and an equivalence test
asserting TOP with a large line limit matches RETR output. All
scripted session tests from the previous commit pass unchanged.

Signed-off-by: James Hillyerd <james@hillyerd.com>

pkg/server/pop3/handler.go

@@ -293,101 +293,21 @@ func (s *Session) transactionHandler(cmd string, args []string) {
 		}
 		s.send(fmt.Sprintf("+OK %v %v", count, size))
 	case "LIST":
-		if len(args) > 1 {
+		s.listMessages("LIST", args, func(m storage.Message) string {
-			s.logger.Warn().Msgf("LIST command had more than 1 argument")
+			return strconv.FormatInt(m.Size(), 10)
-			s.send("-ERR LIST command must have zero or one argument")
+		})
-			return
-		}
-		if len(args) == 1 {
-			msgNum, err := strconv.ParseInt(args[0], 10, 32)
-			if err != nil {
-				s.logger.Warn().Msgf("LIST command argument was not an integer")
-				s.send("-ERR LIST command requires an integer argument")
-				return
-			}
-			if msgNum < 1 {
-				s.logger.Warn().Msgf("LIST command argument was less than 1")
-				s.send("-ERR LIST argument must be greater than 0")
-				return
-			}
-			if int(msgNum) > len(s.messages) {
-				s.logger.Warn().Msgf("LIST command argument was greater than number of messages")
-				s.send("-ERR LIST argument must not exceed the number of messages")
-				return
-			}
-			if !s.retain[msgNum-1] {
-				s.logger.Warn().Msgf("Client tried to LIST a message it had deleted")
-				s.send(fmt.Sprintf("-ERR You deleted message %v", msgNum))
-				return
-			}
-			s.send(fmt.Sprintf("+OK %v %v", msgNum, s.messages[msgNum-1].Size()))
-		} else {
-			s.send(fmt.Sprintf("+OK Listing %v messages", s.msgCount))
-			for i, msg := range s.messages {
-				if s.retain[i] {
-					s.send(fmt.Sprintf("%v %v", i+1, msg.Size()))
-				}
-			}
-			s.send(".")
-		}
 	case "UIDL":
-		if len(args) > 1 {
+		s.listMessages("UIDL", args, func(m storage.Message) string {
-			s.logger.Warn().Msgf("UIDL command had more than 1 argument")
+			return m.ID()
-			s.send("-ERR UIDL command must have zero or one argument")
+		})
-			return
-		}
-		if len(args) == 1 {
-			msgNum, err := strconv.ParseInt(args[0], 10, 32)
-			if err != nil {
-				s.logger.Warn().Msgf("UIDL command argument was not an integer")
-				s.send("-ERR UIDL command requires an integer argument")
-				return
-			}
-			if msgNum < 1 {
-				s.logger.Warn().Msgf("UIDL command argument was less than 1")
-				s.send("-ERR UIDL argument must be greater than 0")
-				return
-			}
-			if int(msgNum) > len(s.messages) {
-				s.logger.Warn().Msgf("UIDL command argument was greater than number of messages")
-				s.send("-ERR UIDL argument must not exceed the number of messages")
-				return
-			}
-			if !s.retain[msgNum-1] {
-				s.logger.Warn().Msgf("Client tried to UIDL a message it had deleted")
-				s.send(fmt.Sprintf("-ERR You deleted message %v", msgNum))
-				return
-			}
-			s.send(fmt.Sprintf("+OK %v %v", msgNum, s.messages[msgNum-1].ID()))
-		} else {
-			s.send(fmt.Sprintf("+OK Listing %v messages", s.msgCount))
-			for i, msg := range s.messages {
-				if s.retain[i] {
-					s.send(fmt.Sprintf("%v %v", i+1, msg.ID()))
-				}
-			}
-			s.send(".")
-		}
 	case "DELE":
 		if len(args) != 1 {
 			s.logger.Warn().Msgf("DELE command had invalid number of arguments")
 			s.send("-ERR DELE command requires a single argument")
 			return
 		}
-		msgNum, err := strconv.ParseInt(args[0], 10, 32)
+		msgNum, ok := s.validateMsgNum("DELE", "DELE", args[0])
-		if err != nil {
+		if !ok {
-			s.logger.Warn().Msgf("DELE command argument was not an integer")
-			s.send("-ERR DELE command requires an integer argument")
-			return
-		}
-		if msgNum < 1 {
-			s.logger.Warn().Msgf("DELE command argument was less than 1")
-			s.send("-ERR DELE argument must be greater than 0")
-			return
-		}
-		if int(msgNum) > len(s.messages) {
-			s.logger.Warn().Msgf("DELE command argument was greater than number of messages")
-			s.send("-ERR DELE argument must not exceed the number of messages")
 			return
 		}
 		if s.retain[msgNum-1] {
@@ -404,49 +324,25 @@ func (s *Session) transactionHandler(cmd string, args []string) {
 			s.send("-ERR RETR command requires a single argument")
 			return
 		}
-		msgNum, err := strconv.ParseInt(args[0], 10, 32)
+		msgNum, ok := s.validateMsgNum("RETR", "RETR", args[0])
-		if err != nil {
+		if !ok {
-			s.logger.Warn().Msgf("RETR command argument was not an integer")
-			s.send("-ERR RETR command requires an integer argument")
-			return
-		}
-		if msgNum < 1 {
-			s.logger.Warn().Msgf("RETR command argument was less than 1")
-			s.send("-ERR RETR argument must be greater than 0")
-			return
-		}
-		if int(msgNum) > len(s.messages) {
-			s.logger.Warn().Msgf("RETR command argument was greater than number of messages")
-			s.send("-ERR RETR argument must not exceed the number of messages")
 			return
 		}
 		s.send(fmt.Sprintf("+OK %v bytes follows", s.messages[msgNum-1].Size()))
-		s.sendMessage(s.messages[msgNum-1])
+		s.sendMessageLines(s.messages[msgNum-1], -1)
 	case "TOP":
 		if len(args) != 2 {
 			s.logger.Warn().Msgf("TOP command had invalid number of arguments")
 			s.send("-ERR TOP command requires two arguments")
 			return
 		}
-		msgNum, err := strconv.ParseInt(args[0], 10, 32)
+		msgNum, ok := s.validateMsgNum("TOP", "TOP first", args[0])
-		if err != nil {
+		if !ok {
-			s.logger.Warn().Msgf("TOP command first argument was not an integer")
-			s.send("-ERR TOP command requires an integer argument")
-			return
-		}
-		if msgNum < 1 {
-			s.logger.Warn().Msgf("TOP command first argument was less than 1")
-			s.send("-ERR TOP first argument must be greater than 0")
-			return
-		}
-		if int(msgNum) > len(s.messages) {
-			s.logger.Warn().Msgf("TOP command first argument was greater than number of messages")
-			s.send("-ERR TOP first argument must not exceed the number of messages")
 			return
 		}
 
 		var lines int64
-		lines, err = strconv.ParseInt(args[1], 10, 32)
+		lines, err := strconv.ParseInt(args[1], 10, 32)
 		if err != nil {
 			s.logger.Warn().Msgf("TOP command second argument was not an integer")
 			s.send("-ERR TOP command requires an integer argument")
@@ -458,7 +354,7 @@ func (s *Session) transactionHandler(cmd string, args []string) {
 			return
 		}
 		s.send("+OK Top of message follows")
-		s.sendMessageTop(s.messages[msgNum-1], int(lines))
+		s.sendMessageLines(s.messages[msgNum-1], int(lines))
 	case "QUIT":
 		s.send("+OK We will process your deletes")
 		s.processDeletes()
@@ -475,41 +371,65 @@ func (s *Session) transactionHandler(cmd string, args []string) {
 	}
 }
 
-// Send the contents of the message to the client
+// listMessages implements the LIST and UIDL commands, which differ only in
-func (s *Session) sendMessage(msg storage.Message) {
+// the value reported for each message.
-	reader, err := msg.Source()
+func (s *Session) listMessages(cmd string, args []string, value func(m storage.Message) string) {
-	if err != nil {
+	if len(args) > 1 {
-		s.logger.Error().Msgf("Failed to read message for RETR command")
+		s.logger.Warn().Msgf("%v command had more than 1 argument", cmd)
-		s.send("-ERR Failed to RETR that message, internal error")
+		s.send(fmt.Sprintf("-ERR %v command must have zero or one argument", cmd))
 		return
 	}
-	defer func() {
+	if len(args) == 1 {
-		if err := reader.Close(); err != nil {
+		msgNum, ok := s.validateMsgNum(cmd, cmd, args[0])
-			s.logger.Error().Msgf("Failed to close message: %v", err)
+		if !ok {
+			return
 		}
-	}()
+		if !s.retain[msgNum-1] {
-
+			s.logger.Warn().Msgf("Client tried to %v a message it had deleted", cmd)
-	scanner := bufio.NewScanner(reader)
+			s.send(fmt.Sprintf("-ERR You deleted message %v", msgNum))
-	for scanner.Scan() {
+			return
-		line := scanner.Text()
-		// Lines starting with . must be prefixed with another .
-		if strings.HasPrefix(line, ".") {
-			line = "." + line
 		}
-		s.send(line)
+		s.send(fmt.Sprintf("+OK %v %v", msgNum, value(s.messages[msgNum-1])))
-	}
-
-	if err = scanner.Err(); err != nil {
-		s.logger.Error().Msgf("Failed to read message for RETR command")
-		s.send(".")
-		s.send("-ERR Failed to RETR that message, internal error")
 		return
 	}
+	s.send(fmt.Sprintf("+OK Listing %v messages", s.msgCount))
+	for i, msg := range s.messages {
+		if s.retain[i] {
+			s.send(fmt.Sprintf("%v %v", i+1, value(msg)))
+		}
+	}
 	s.send(".")
 }
 
-// Send the headers plus the top N lines to the client
+// validateMsgNum parses arg as a message number and range-checks it against
-func (s *Session) sendMessageTop(msg storage.Message, lineCount int) {
+// the messages in the mailbox, sending a -ERR response and returning ok=false
+// when invalid. cmd labels the command in responses; argLabel labels the
+// argument (e.g. "RETR", "TOP first").
+func (s *Session) validateMsgNum(cmd, argLabel, arg string) (msgNum int, ok bool) {
+	n, err := strconv.ParseInt(arg, 10, 32)
+	if err != nil {
+		s.logger.Warn().Msgf("%v command argument was not an integer", cmd)
+		s.send(fmt.Sprintf("-ERR %v command requires an integer argument", cmd))
+		return 0, false
+	}
+	if n < 1 {
+		s.logger.Warn().Msgf("%v argument was less than 1", argLabel)
+		s.send(fmt.Sprintf("-ERR %v argument must be greater than 0", argLabel))
+		return 0, false
+	}
+	if int(n) > len(s.messages) {
+		s.logger.Warn().Msgf("%v argument was greater than number of messages", argLabel)
+		s.send(fmt.Sprintf("-ERR %v argument must not exceed the number of messages", argLabel))
+		return 0, false
+	}
+	return int(n), true
+}
+
+// sendMessageLines sends the raw contents of msg to the client, dot-stuffing
+// any lines that begin with a period. If maxBodyLines is negative the entire
+// message is sent; otherwise the headers (and separating blank line) plus at
+// most maxBodyLines lines of the body are sent.
+func (s *Session) sendMessageLines(msg storage.Message, maxBodyLines int) {
 	reader, err := msg.Source()
 	if err != nil {
 		s.logger.Error().Msgf("Failed to read message for RETR command")
@@ -531,17 +451,13 @@ func (s *Session) sendMessageTop(msg storage.Message, lineCount int) {
 			line = "." + line
 		}
 		if inBody {
-			// Check if we need to send anymore lines
+			if maxBodyLines >= 0 && maxBodyLines < 1 {
-			if lineCount < 1 {
 				break
-			} else {
-				lineCount--
-			}
-		} else {
-			if line == "" {
-				// We've hit the end of the header
-				inBody = true
 			}
+			maxBodyLines--
+		} else if line == "" {
+			// We've hit the end of the header
+			inBody = true
 		}
 		s.send(line)
 	}

pkg/server/pop3/transaction_test.go

@@ -1,6 +1,7 @@
 package pop3
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"fmt"
@@ -101,8 +102,8 @@ func newMemStore(t *testing.T) storage.Store {
 }
 
 // deliverRaw stores a message with the exact raw content provided, returning
-// its ID and size.
+// its size.
-func deliverRaw(t *testing.T, store storage.Store, mailbox, raw string) (string, int64) {
+func deliverRaw(t *testing.T, store storage.Store, mailbox, raw string) int64 {
 	t.Helper()
 
 	delivery := &message.Delivery{
@@ -112,9 +113,9 @@ func deliverRaw(t *testing.T, store storage.Store, mailbox, raw string) (string,
 		},
 		Reader: io.NopCloser(strings.NewReader(raw)),
 	}
-	id, err := store.AddMessage(delivery)
+	_, err := store.AddMessage(delivery)
 	require.NoError(t, err)
-	return id, int64(len(raw))
+	return int64(len(raw))
 }
 
 // fakeMsg is a storage.Message with canned metadata, used to exercise RETR/TOP
@@ -532,7 +533,7 @@ func TestTransactionTopDeletedMessage(t *testing.T) {
 
 func TestTransactionRetrDotStuffing(t *testing.T) {
 	ds := newMemStore(t)
-	_, size := deliverRaw(t, ds, "mailbox", "Subject: dots\r\n\r\n.leading\r\n..double\r\n.\r\nplain\r\n")
+	size := deliverRaw(t, ds, "mailbox", "Subject: dots\r\n\r\n.leading\r\n..double\r\n.\r\nplain\r\n")
 	server := setupPOPServer(t, ds, false, false)
 
 	script := slices.Concat(loginSteps(1), []scriptStep{
@@ -764,3 +765,79 @@ func TestNewServerCertificateError(t *testing.T) {
 	}, test.NewStore())
 	require.Error(t, err)
 }
+
+// TestValidateMsgNum exercises the shared message-number validation helper
+// directly.
+func TestValidateMsgNum(t *testing.T) {
+	server := setupPOPServer(t, newMemStore(t), false, false)
+
+	serverConn, clientConn := net.Pipe()
+	t.Cleanup(func() { _ = clientConn.Close() })
+
+	// Drain server->client output so sends do not block.
+	output := make(chan string, 16)
+	go func() {
+		sc := bufio.NewScanner(clientConn)
+		for sc.Scan() {
+			output <- sc.Text()
+		}
+		close(output)
+	}()
+
+	ssn := NewSession(server, 1, &mockConn{serverConn}, zerolog.New(zerolog.NewTestWriter(t)))
+	ssn.messages = make([]storage.Message, 3)
+	ssn.retainAll()
+
+	tcs := []struct {
+		name     string
+		arg      string
+		wantNum  int
+		wantOK   bool
+		wantSend string
+	}{
+		{name: "first message", arg: "1", wantNum: 1, wantOK: true},
+		{name: "last message", arg: "3", wantNum: 3, wantOK: true},
+		{name: "not an integer", arg: "x", wantSend: "-ERR TEST command requires an integer argument"},
+		{name: "integer overflow", arg: "9999999999999", wantSend: "-ERR TEST command requires an integer argument"},
+		{name: "zero", arg: "0", wantSend: "-ERR TEST argument must be greater than 0"},
+		{name: "negative", arg: "-1", wantSend: "-ERR TEST argument must be greater than 0"},
+		{name: "exceeds message count", arg: "4", wantSend: "-ERR TEST argument must not exceed the number of messages"},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			num, ok := ssn.validateMsgNum("TEST", "TEST", tc.arg)
+			assert.Equal(t, tc.wantOK, ok)
+			assert.Equal(t, tc.wantNum, num)
+
+			if tc.wantSend != "" {
+				select {
+				case line := <-output:
+					assert.Equal(t, tc.wantSend, line)
+				case <-time.After(time.Second):
+					t.Error("expected an error response")
+				}
+			}
+		})
+	}
+}
+
+// TestTopUnlimitedLinesMatchesRetr guards the shared implementation of RETR
+// and TOP: a TOP line limit larger than the message must produce the same
+// output as RETR.
+func TestTopUnlimitedLinesMatchesRetr(t *testing.T) {
+	ds := newMemStore(t)
+	size := deliverRaw(t, ds, "mailbox", "Subject: equivalence\r\n\r\nline 1\r\nline 2\r\n.leading\r\n")
+	server := setupPOPServer(t, ds, false, false)
+
+	body := []string{"Subject: equivalence", "", "line 1", "line 2", "..leading", "."}
+	retrExpect := slices.Concat([]string{fmt.Sprintf("+OK %v bytes follows", size)}, body)
+	topExpect := slices.Concat([]string{"+OK Top of message follows"}, body)
+
+	script := slices.Concat(loginSteps(1), []scriptStep{
+		{"RETR 1", retrExpect},
+		{"TOP 1 999999", topExpect},
+		{"QUIT", []string{"+OK We will process your deletes"}},
+	})
+	playPOP3Session(t, server, script)
+}