mox/imapclient/client.go

/*
Package imapclient provides an IMAP4 client, primarily for testing the IMAP4 server.

Commands can be sent to the server free-form, but responses are parsed strictly.
Behaviour that may not be required by the IMAP4 specification may be expected by
this client.
*/
package imapclient

/*
- Try to keep the parsing method names and the types similar to the ABNF names in the RFCs.

- todo: have mode for imap4rev1 vs imap4rev2, refusing what is not allowed. we are accepting too much now.
- todo: stricter parsing. xnonspace() and xword() should be replaced by proper parsers.
*/

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"log/slog"
	"net"
	"reflect"
	"strings"

	"github.com/mjl-/mox/mlog"
	"github.com/mjl-/mox/moxio"
)

// Conn is an IMAP connection to a server.
type Conn struct {
	// Connection, may be original TCP or TLS connection. Reads go through c.br, and
	// writes through c.xbw. The "x" for the writes indicate that failed writes cause
	// an i/o panic, which is either turned into a returned error, or passed on (see
	// boolean panic). The reader and writer wrap a tracing reading/writer and may wrap
	// flate compression.
	conn         net.Conn
	connBroken   bool // If connection is broken, we won't flush (and write) again.
	br           *bufio.Reader
	tr           *moxio.TraceReader
	xbw          *bufio.Writer
	compress     bool // If compression is enabled, we must flush flateWriter and its target original bufio writer.
	xflateWriter *moxio.FlateWriter
	xflateBW     *bufio.Writer
	xtw          *moxio.TraceWriter

	log       mlog.Log
	errHandle func(err error) // If set, called for all errors. Can panic. Used for imapserver tests.
	tagGen    int
	record    bool // If true, bytes read are added to recordBuf. recorded() resets.
	recordBuf []byte

	Preauth      bool
	LastTag      string
	CapAvailable map[Capability]struct{} // Capabilities available at server, from CAPABILITY command or response code. All uppercase.
	CapEnabled   map[Capability]struct{} // Capabilities enabled through ENABLE command. All uppercase.
}

// Error is a parse or other protocol error.
type Error struct{ err error }

func (e Error) Error() string {
	return e.err.Error()
}

func (e Error) Unwrap() error {
	return e.err
}

// Opts has optional fields that influence behaviour of a Conn.
type Opts struct {
	Logger *slog.Logger

	// Error is called for both IMAP-level and connection-level errors. Is allowed to
	// call panic.
	Error func(err error)
}

// New creates a new client on conn.
//
// The initial untagged greeting response is read and must be "OK" or
// "PREAUTH". If preauth, the connection is already in authenticated state,
// typically through TLS client certificate. This is indicated in Conn.Preauth.
//
// Logging is written to log, in particular IMAP protocol traces are written with
// prefixes "CR: " and "CW: " (client read/write) as quoted strings at levels
// Debug-4, with authentication messages at Debug-6 and (user) data at level
// Debug-8.
func New(conn net.Conn, opts *Opts) (client *Conn, rerr error) {
	c := Conn{
		conn:         conn,
		CapAvailable: map[Capability]struct{}{},
		CapEnabled:   map[Capability]struct{}{},
	}

	var clog *slog.Logger
	if opts != nil {
		c.errHandle = opts.Error
		clog = opts.Logger
	} else {
		clog = slog.Default()
	}
	c.log = mlog.New("imapclient", clog)

	c.tr = moxio.NewTraceReader(c.log, "CR: ", &c)
	c.br = bufio.NewReader(c.tr)

	// Writes are buffered and write to Conn, which may panic.
	c.xtw = moxio.NewTraceWriter(c.log, "CW: ", &c)
	c.xbw = bufio.NewWriter(c.xtw)

	defer c.recover(&rerr)
	tag := c.xnonspace()
	if tag != "*" {
		c.xerrorf("expected untagged *, got %q", tag)
	}
	c.xspace()
	ut := c.xuntagged()
	switch x := ut.(type) {
	case UntaggedResult:
		if x.Status != OK {
			c.xerrorf("greeting, got status %q, expected OK", x.Status)
		}
		return &c, nil
	case UntaggedPreauth:
		c.Preauth = true
		return &c, nil
	case UntaggedBye:
		c.xerrorf("greeting: server sent bye")
	default:
		c.xerrorf("unexpected untagged %v", ut)
	}
	panic("not reached")
}

func (c *Conn) recover(rerr *error) {
	x := recover()
	if x == nil {
		return
	}
	err, ok := x.(Error)
	if !ok {
		panic(x)
	}
	if c.errHandle != nil {
		c.errHandle(err)
	}
	*rerr = err
}

func (c *Conn) xerrorf(format string, args ...any) {
	panic(Error{fmt.Errorf(format, args...)})
}

func (c *Conn) xcheckf(err error, format string, args ...any) {
	if err != nil {
		c.xerrorf("%s: %w", fmt.Sprintf(format, args...), err)
	}
}

func (c *Conn) xcheck(err error) {
	if err != nil {
		panic(err)
	}
}

// Write writes directly to underlying connection (TCP, TLS). For internal use
// only, to implement io.Writer. Write errors do take the connection's panic mode
// into account, i.e. Write can panic.
func (c *Conn) Write(buf []byte) (n int, rerr error) {
	defer c.recover(&rerr)

	n, rerr = c.conn.Write(buf)
	if rerr != nil {
		c.connBroken = true
	}
	c.xcheckf(rerr, "write")
	return n, nil
}

// Read reads directly from the underlying connection (TCP, TLS). For internal use
// only, to implement io.Reader.
func (c *Conn) Read(buf []byte) (n int, err error) {
	return c.conn.Read(buf)
}

func (c *Conn) xflush() {
	// Not writing any more when connection is broken.
	if c.connBroken {
		return
	}

	err := c.xbw.Flush()
	c.xcheckf(err, "flush")

	// If compression is active, we need to flush the deflate stream.
	if c.compress {
		err := c.xflateWriter.Flush()
		c.xcheckf(err, "flush deflate")
		err = c.xflateBW.Flush()
		c.xcheckf(err, "flush deflate buffer")
	}
}

func (c *Conn) xtraceread(level slog.Level) func() {
	c.tr.SetTrace(level)
	return func() {
		c.tr.SetTrace(mlog.LevelTrace)
	}
}

func (c *Conn) xtracewrite(level slog.Level) func() {
	c.xflush()
	c.xtw.SetTrace(level)
	return func() {
		c.xflush()
		c.xtw.SetTrace(mlog.LevelTrace)
	}
}

// Close closes the connection, flushing and closing any compression and TLS layer.
//
// You may want to call Logout first. Closing a connection with a mailbox with
// deleted messages not yet expunged will not expunge those messages.
//
// Closing a TLS connection that is logged out, or closing a TLS connection with
// compression enabled (i.e. two layered streams), may cause spurious errors
// because the server may immediate close the underlying connection when it sees
// the connection is being closed.
func (c *Conn) Close() (rerr error) {
	defer c.recover(&rerr)

	if c.conn == nil {
		return nil
	}
	if !c.connBroken && c.xflateWriter != nil {
		err := c.xflateWriter.Close()
		c.xcheckf(err, "close deflate writer")
		err = c.xflateBW.Flush()
		c.xcheckf(err, "flush deflate buffer")
		c.xflateWriter = nil
		c.xflateBW = nil
	}
	err := c.conn.Close()
	c.xcheckf(err, "close connection")
	c.conn = nil
	return
}

// TLSConnectionState returns the TLS connection state if the connection uses TLS.
func (c *Conn) TLSConnectionState() *tls.ConnectionState {
	if conn, ok := c.conn.(*tls.Conn); ok {
		cs := conn.ConnectionState()
		return &cs
	}
	return nil
}

// Commandf writes a free-form IMAP command to the server. An ending \r\n is
// written too.
// If tag is empty, a next unique tag is assigned.
func (c *Conn) Commandf(tag string, format string, args ...any) (rerr error) {
	defer c.recover(&rerr)

	if tag == "" {
		tag = c.nextTag()
	}
	c.LastTag = tag

	fmt.Fprintf(c.xbw, "%s %s\r\n", tag, fmt.Sprintf(format, args...))
	c.xflush()
	return
}

func (c *Conn) nextTag() string {
	c.tagGen++
	return fmt.Sprintf("x%03d", c.tagGen)
}

// Response reads from the IMAP server until a tagged response line is found.
// The tag must be the same as the tag for the last written command.
// Result holds the status of the command. The caller must check if this the status is OK.
func (c *Conn) Response() (untagged []Untagged, result Result, rerr error) {
	defer c.recover(&rerr)

	for {
		tag := c.xnonspace()
		c.xspace()
		if tag == "*" {
			untagged = append(untagged, c.xuntagged())
			continue
		}

		if tag != c.LastTag {
			c.xerrorf("got tag %q, expected %q", tag, c.LastTag)
		}

		status := c.xstatus()
		c.xspace()
		result = c.xresult(status)
		c.xcrlf()
		return
	}
}

// ReadUntagged reads a single untagged response line.
// Useful for reading lines from IDLE.
func (c *Conn) ReadUntagged() (untagged Untagged, rerr error) {
	defer c.recover(&rerr)

	tag := c.xnonspace()
	if tag != "*" {
		c.xerrorf("got tag %q, expected untagged", tag)
	}
	c.xspace()
	ut := c.xuntagged()
	return ut, nil
}

// Readline reads a line, including CRLF.
// Used with IDLE and synchronous literals.
func (c *Conn) Readline() (line string, rerr error) {
	defer c.recover(&rerr)

	line, err := c.br.ReadString('\n')
	c.xcheckf(err, "read line")
	return line, nil
}

// ReadContinuation reads a line. If it is a continuation, i.e. starts with a +, it
// is returned without leading "+ " and without trailing crlf. Otherwise, a command
// response is returned. A successfully read continuation can return an empty line.
// Callers should check rerr and result.Status being empty to check if a
// continuation was read.
func (c *Conn) ReadContinuation() (line string, untagged []Untagged, result Result, rerr error) {
	defer c.recover(&rerr)

	if !c.peek('+') {
		untagged, result, rerr = c.Response()
		if result.Status == OK {
			c.xerrorf("unexpected OK instead of continuation")
		}
		return
	}
	c.xtake("+ ")
	line, err := c.Readline()
	c.xcheckf(err, "read line")
	line = strings.TrimSuffix(line, "\r\n")
	return
}

// Writelinef writes the formatted format and args as a single line, adding CRLF.
// Used with IDLE and synchronous literals.
func (c *Conn) Writelinef(format string, args ...any) (rerr error) {
	defer c.recover(&rerr)

	s := fmt.Sprintf(format, args...)
	fmt.Fprintf(c.xbw, "%s\r\n", s)
	c.xflush()
	return nil
}

// WriteSyncLiteral first writes the synchronous literal size, then reads the
// continuation "+" and finally writes the data.
func (c *Conn) WriteSyncLiteral(s string) (untagged []Untagged, rerr error) {
	defer c.recover(&rerr)

	fmt.Fprintf(c.xbw, "{%d}\r\n", len(s))
	c.xflush()

	plus, err := c.br.Peek(1)
	c.xcheckf(err, "read continuation")
	if plus[0] == '+' {
		_, err = c.Readline()
		c.xcheckf(err, "read continuation line")

		defer c.xtracewrite(mlog.LevelTracedata)()
		_, err = c.xbw.Write([]byte(s))
		c.xcheckf(err, "write literal data")
		c.xtracewrite(mlog.LevelTrace)
		return nil, nil
	}
	untagged, result, err := c.Response()
	if err == nil && result.Status == OK {
		c.xerrorf("no continuation, but invalid ok response (%q)", result.More)
	}
	return untagged, fmt.Errorf("no continuation (%s)", result.Status)
}

// Transactf writes format and args as an IMAP command, using Commandf with an
// empty tag. I.e. format must not contain a tag. Transactf then reads a response
// using ReadResponse and checks the result status is OK.
func (c *Conn) Transactf(format string, args ...any) (untagged []Untagged, result Result, rerr error) {
	defer c.recover(&rerr)

	err := c.Commandf("", format, args...)
	if err != nil {
		return nil, Result{}, err
	}
	return c.ResponseOK()
}

func (c *Conn) ResponseOK() (untagged []Untagged, result Result, rerr error) {
	untagged, result, rerr = c.Response()
	if rerr != nil {
		return nil, Result{}, rerr
	}
	if result.Status != OK {
		c.xerrorf("response status %q, expected OK", result.Status)
	}
	return untagged, result, rerr
}

func (c *Conn) xgetUntagged(l []Untagged, dst any) {
	if len(l) != 1 {
		c.xerrorf("got %d untagged, expected 1: %v", len(l), l)
	}
	got := l[0]
	gotv := reflect.ValueOf(got)
	dstv := reflect.ValueOf(dst)
	if gotv.Type() != dstv.Type().Elem() {
		c.xerrorf("got %v, expected %v", gotv.Type(), dstv.Type().Elem())
	}
	dstv.Elem().Set(gotv)
}