gomail/internal/managesieve/session.go

package managesieve

import (
	"bufio"
	"context"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"log/slog"
	"net"
	"strconv"
	"strings"
	"time"

	"gomail/internal/auth"
	"gomail/internal/db"
	"gomail/internal/sieve"
	"github.com/google/uuid"
)

type session struct {
	conn      net.Conn
	rw        *bufio.ReadWriter
	server    *Server
	tlsActive bool
	user      *db.User
}

func newSession(conn net.Conn, server *Server) *session {
	_, isTLS := conn.(*tls.Conn)
	return &session{
		conn:      conn,
		rw:        bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)),
		server:    server,
		tlsActive: isTLS,
	}
}

func (s *session) run(ctx context.Context) {
	s.sendCapabilities()

	for {
		select {
		case <-ctx.Done():
			s.writeLine(`BYE "server shutting down"`)
			return
		default:
		}

		s.conn.SetReadDeadline(time.Now().Add(idleTimeout))
		line, err := s.readLine()
		if err != nil {
			if err != io.EOF {
				slog.Debug("ManageSieve read error", "err", err)
			}
			return
		}
		if !s.dispatch(line) {
			return
		}
	}
}

func (s *session) sendCapabilities() {
	s.writeLine(`"IMPLEMENTATION" "GoMail ManageSieve"`)
	s.writeLine(`"SIEVE" "fileinto"`)
	s.writeLine(`"VERSION" "1.0"`)
	if !s.tlsActive {
		s.writeLine(`"STARTTLS"`)
	}
	s.writeLine("OK")
}

func (s *session) dispatch(line string) bool {
	verb, rest := splitVerb(line)
	switch strings.ToUpper(verb) {
	case "CAPABILITY":
		s.sendCapabilities()
	case "STARTTLS":
		s.cmdStartTLS()
	case "AUTHENTICATE":
		s.cmdAuthenticate(rest)
	case "LOGOUT":
		s.writeLine("OK")
		return false
	case "PUTSCRIPT":
		s.cmdPutScript(rest)
	case "GETSCRIPT":
		s.cmdGetScript(rest)
	case "LISTSCRIPTS":
		s.cmdListScripts()
	case "SETACTIVE":
		s.cmdSetActive(rest)
	case "DELETESCRIPT":
		s.cmdDeleteScript(rest)
	case "NOOP":
		s.writeLine("OK")
	default:
		s.writeLine(`NO "command not recognized"`)
	}
	return true
}

func (s *session) cmdStartTLS() {
	if s.tlsActive {
		s.writeLine(`NO "TLS already active"`)
		return
	}
	s.writeLine("OK")
	tlsConn := tls.Server(s.conn, s.server.tlsConf)
	if err := tlsConn.HandshakeContext(context.Background()); err != nil {
		return
	}
	s.conn = tlsConn
	s.rw = bufio.NewReadWriter(bufio.NewReader(tlsConn), bufio.NewWriter(tlsConn))
	s.tlsActive = true
}

// cmdAuthenticate handles AUTHENTICATE "PLAIN" <base64> — the SASL PLAIN
// mechanism, same as SMTP/IMAP's AUTH PLAIN, adapted to ManageSieve's quoted
// string argument syntax rather than a bare base64 token.
func (s *session) cmdAuthenticate(rest string) {
	if !s.tlsActive {
		s.writeLine(`NO "authentication requires TLS — use STARTTLS first"`)
		return
	}

	parts := splitQuotedArgs(rest)
	if len(parts) < 1 || strings.ToUpper(strings.Trim(parts[0], `"`)) != "PLAIN" {
		s.writeLine(`NO "only AUTHENTICATE PLAIN is supported"`)
		return
	}

	var b64 string
	if len(parts) >= 2 {
		b64 = strings.Trim(parts[1], `"`)
	} else {
		s.writeLine("{0}")
		line, err := s.readLine()
		if err != nil {
			return
		}
		b64 = line
	}

	decoded, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		s.writeLine(`NO "malformed SASL response"`)
		return
	}
	fields := strings.SplitN(string(decoded), "\x00", 3)
	if len(fields) != 3 {
		s.writeLine(`NO "malformed SASL PLAIN payload"`)
		return
	}
	username, password := fields[1], fields[2]

	user, ok := auth.Authenticate(s.server.database, username, password, auth.ScopeIMAP)
	if !ok {
		s.writeLine(`NO "authentication failed"`)
		return
	}
	s.user = user
	s.writeLine("OK")
}

func (s *session) requireAuth() bool {
	if s.user == nil {
		s.writeLine(`NO "authentication required"`)
		return false
	}
	return true
}

// cmdPutScript handles: PUTSCRIPT "name" {N+}\r\n<N bytes of script>\r\n
func (s *session) cmdPutScript(rest string) {
	if !s.requireAuth() {
		return
	}
	parts := splitQuotedArgs(rest)
	if len(parts) < 1 {
		s.writeLine(`NO "PUTSCRIPT requires a script name"`)
		return
	}
	name := strings.Trim(parts[0], `"`)

	scriptText, err := s.readLiteralFromRemainder(rest)
	if err != nil {
		s.writeLine(`NO "expected script literal: ` + err.Error() + `"`)
		return
	}

	if _, err := sieve.Parse(scriptText); err != nil {
		s.writeLine(`NO "script failed to parse: ` + escapeQuoted(err.Error()) + `"`)
		return
	}

	if err := s.server.database.UpsertSieveScript(&db.SieveScript{
		ID: uuid.NewString(), UserID: s.user.ID, Name: name, ScriptText: scriptText,
	}); err != nil {
		s.writeLine(`NO "storage error"`)
		return
	}
	s.writeLine("OK")
}

func (s *session) cmdGetScript(rest string) {
	if !s.requireAuth() {
		return
	}
	name := strings.Trim(strings.TrimSpace(rest), `"`)
	script, err := s.server.database.GetSieveScript(s.user.ID, name)
	if err != nil {
		s.writeLine(`NO "script not found"`)
		return
	}
	s.writeLine(fmt.Sprintf("{%d}", len(script.ScriptText)))
	s.rw.WriteString(script.ScriptText)
	s.rw.WriteString("\r\n")
	s.rw.Flush()
	s.writeLine("OK")
}

func (s *session) cmdListScripts() {
	if !s.requireAuth() {
		return
	}
	scripts, err := s.server.database.ListSieveScripts(s.user.ID)
	if err != nil {
		s.writeLine(`NO "storage error"`)
		return
	}
	for _, sc := range scripts {
		if sc.Active {
			s.writeLine(fmt.Sprintf(`"%s" ACTIVE`, sc.Name))
		} else {
			s.writeLine(fmt.Sprintf(`"%s"`, sc.Name))
		}
	}
	s.writeLine("OK")
}

func (s *session) cmdSetActive(rest string) {
	if !s.requireAuth() {
		return
	}
	name := strings.Trim(strings.TrimSpace(rest), `"`)
	if name == "" {
		// Empty name deactivates all scripts, per RFC 5804 §2.9.
		s.server.database.Exec(`UPDATE sieve_scripts SET active = 0 WHERE user_id = ?`, s.user.ID)
		s.writeLine("OK")
		return
	}
	if err := s.server.database.SetActiveSieveScript(s.user.ID, name); err != nil {
		s.writeLine(`NO "script not found"`)
		return
	}
	s.writeLine("OK")
}

func (s *session) cmdDeleteScript(rest string) {
	if !s.requireAuth() {
		return
	}
	name := strings.Trim(strings.TrimSpace(rest), `"`)
	if err := s.server.database.DeleteSieveScript(s.user.ID, name); err != nil {
		s.writeLine(`NO "delete failed"`)
		return
	}
	s.writeLine("OK")
}

// ── I/O helpers ─────────────────────────────────────────────────────────────────

func (s *session) writeLine(line string) {
	s.rw.WriteString(line + "\r\n")
	s.rw.Flush()
}

func (s *session) readLine() (string, error) {
	line, err := s.rw.ReadString('\n')
	if err != nil {
		return "", err
	}
	return strings.TrimRight(line, "\r\n"), nil
}

// readLiteralFromRemainder expects the command line's remainder to end in a
// {N} or {N+} literal announcement (RFC 5804 reuses IMAP-style literal
// syntax) and reads exactly N raw bytes following it.
func (s *session) readLiteralFromRemainder(rest string) (string, error) {
	idx := strings.LastIndex(rest, "{")
	if idx == -1 || !strings.HasSuffix(strings.TrimSpace(rest), "}") {
		return "", fmt.Errorf("no literal size announced")
	}
	sizeStr := strings.TrimSuffix(strings.TrimSpace(rest[idx+1:]), "}")
	sizeStr = strings.TrimSuffix(sizeStr, "+")
	n, err := strconv.Atoi(sizeStr)
	if err != nil {
		return "", fmt.Errorf("invalid literal size: %w", err)
	}

	buf := make([]byte, n)
	if _, err := io.ReadFull(s.rw, buf); err != nil {
		return "", fmt.Errorf("reading literal: %w", err)
	}
	s.rw.ReadString('\n') // consume trailing CRLF after the literal bytes
	return string(buf), nil
}

func splitVerb(line string) (verb, rest string) {
	line = strings.TrimSpace(line)
	i := strings.IndexAny(line, " \t")
	if i < 0 {
		return line, ""
	}
	return line[:i], strings.TrimSpace(line[i+1:])
}

// splitQuotedArgs splits `"arg1" "arg2"` into ["arg1","arg2"] (quotes kept,
// stripped by callers as needed) — tolerant of a trailing {N+} literal
// marker, which callers handle separately via readLiteralFromRemainder.
func splitQuotedArgs(s string) []string {
	var args []string
	i := 0
	for i < len(s) {
		for i < len(s) && s[i] == ' ' {
			i++
		}
		if i >= len(s) {
			break
		}
		if s[i] == '"' {
			j := i + 1
			for j < len(s) && s[j] != '"' {
				j++
			}
			if j < len(s) {
				args = append(args, s[i:j+1])
				i = j + 1
			} else {
				break
			}
		} else {
			j := i
			for j < len(s) && s[j] != ' ' {
				j++
			}
			args = append(args, s[i:j])
			i = j
		}
	}
	return args
}

func escapeQuoted(s string) string {
	return strings.ReplaceAll(s, `"`, `'`)
}