imapclient: add COMPRESS

capability.go

@@ -148,6 +148,19 @@ func (set CapSet) AuthMechanisms() []string {
 	return l
 }
 
+// CompressAlgorithms returns the list of supported compresson mechanisms.
+func (set CapSet) CompressAlgorithms() []string {
+	var l []string
+	for c := range set {
+		if !strings.HasPrefix(string(c), "COMPRESS=") {
+			continue
+		}
+		algo := strings.TrimPrefix(string(c), "COMPRESS=")
+		l = append(l, algo)
+	}
+	return l
+}
+
 // AppendLimit checks the APPENDLIMIT capability.
 //
 // If the server supports APPENDLIMIT, ok is true. If the server doesn't have

imapclient/client.go

@@ -132,6 +132,7 @@ func (options *Options) tlsConfig() *tls.Config {
 // Authenticate, Idle) block the client during their execution.
 type Client struct {
 	conn     net.Conn
+	tlsConn  *tls.Conn
 	options  Options
 	br       *bufio.Reader
 	bw       *bufio.Writer
@@ -617,10 +618,11 @@ func (c *Client) readResponse() error {
 		token    string
 		err      error
 		startTLS *startTLSCommand
+		compress *compressCommand
 	)
 	if tag != "" {
 		token = "response-tagged"
-		startTLS, err = c.readResponseTagged(tag, typ)
+		startTLS, compress, err = c.readResponseTagged(tag, typ)
 	} else {
 		token = "response-data"
 		err = c.readResponseData(typ)
@@ -636,6 +638,9 @@ func (c *Client) readResponse() error {
 	if startTLS != nil {
 		c.upgradeStartTLS(startTLS)
 	}
+	if compress != nil {
+		c.upgradeCompress(compress)
+	}
 
 	return nil
 }
@@ -665,10 +670,10 @@ func (c *Client) readContinueReq() error {
 	return nil
 }
 
-func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand, err error) {
+func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand, compress *compressCommand, err error) {
 	cmd := c.deletePendingCmdByTag(tag)
 	if cmd == nil {
-		return nil, fmt.Errorf("received tagged response with unknown tag %q", tag)
+		return nil, nil, fmt.Errorf("received tagged response with unknown tag %q", tag)
 	}
 
 	// We've removed the command from the pending queue above. Make sure we
@@ -686,14 +691,14 @@ func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand,
 	var code string
 	if hasSP && c.dec.Special('[') { // resp-text-code
 		if !c.dec.ExpectAtom(&code) {
-			return nil, fmt.Errorf("in resp-text-code: %v", c.dec.Err())
+			return nil, nil, fmt.Errorf("in resp-text-code: %v", c.dec.Err())
 		}
 		// TODO: LONGENTRIES and MAXSIZE from METADATA
 		switch code {
 		case "CAPABILITY": // capability-data
 			caps, err := readCapabilities(c.dec)
 			if err != nil {
-				return nil, fmt.Errorf("in capability-data: %v", err)
+				return nil, nil, fmt.Errorf("in capability-data: %v", err)
 			}
 			c.setCaps(caps)
 		case "APPENDUID":
@@ -702,19 +707,19 @@ func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand,
 				uid         imap.UID
 			)
 			if !c.dec.ExpectSP() || !c.dec.ExpectNumber(&uidValidity) || !c.dec.ExpectSP() || !c.dec.ExpectUID(&uid) {
-				return nil, fmt.Errorf("in resp-code-apnd: %v", c.dec.Err())
+				return nil, nil, fmt.Errorf("in resp-code-apnd: %v", c.dec.Err())
 			}
 			if cmd, ok := cmd.(*AppendCommand); ok {
 				cmd.data.UID = uid
 				cmd.data.UIDValidity = uidValidity
 			}
 		case "COPYUID":
 			if !c.dec.ExpectSP() {
-				return nil, c.dec.Err()
+				return nil, nil, c.dec.Err()
 			}
 			uidValidity, srcUIDs, dstUIDs, err := readRespCodeCopyUID(c.dec)
 			if err != nil {
-				return nil, fmt.Errorf("in resp-code-copy: %v", err)
+				return nil, nil, fmt.Errorf("in resp-code-copy: %v", err)
 			}
 			if cmd, ok := cmd.(*CopyCommand); ok {
 				cmd.data.UIDValidity = uidValidity
@@ -727,13 +732,13 @@ func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand,
 			}
 		}
 		if !c.dec.ExpectSpecial(']') {
-			return nil, fmt.Errorf("in resp-text: %v", c.dec.Err())
+			return nil, nil, fmt.Errorf("in resp-text: %v", c.dec.Err())
 		}
 		hasSP = c.dec.SP()
 	}
 	var text string
 	if hasSP && !c.dec.ExpectText(&text) {
-		return nil, fmt.Errorf("in resp-text: %v", c.dec.Err())
+		return nil, nil, fmt.Errorf("in resp-text: %v", c.dec.Err())
 	}
 
 	var cmdErr error
@@ -747,14 +752,17 @@ func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand,
 			Text: text,
 		}
 	default:
-		return nil, fmt.Errorf("in resp-cond-state: expected OK, NO or BAD status condition, but got %v", typ)
+		return nil, nil, fmt.Errorf("in resp-cond-state: expected OK, NO or BAD status condition, but got %v", typ)
 	}
 
 	c.completeCommand(cmd, cmdErr)
 
 	if cmd, ok := cmd.(*startTLSCommand); ok && cmdErr == nil {
 		startTLS = cmd
 	}
+	if cmd, ok := cmd.(*compressCommand); ok && cmdErr == nil {
+		compress = cmd
+	}
 
 	if cmdErr == nil && code != "CAPABILITY" {
 		switch cmd.(type) {
@@ -764,7 +772,7 @@ func (c *Client) readResponseTagged(tag, typ string) (startTLS *startTLSCommand,
 		}
 	}
 
-	return startTLS, nil
+	return startTLS, compress, nil
 }
 
 func (c *Client) readResponseData(typ string) error {

imapclient/compress.go

@@ -0,0 +1,83 @@
+package imapclient
+
+import (
+	"bufio"
+	"bytes"
+	"compress/flate"
+	"io"
+)
+
+// CompressOptions contains options for Client.Compress.
+type CompressOptions struct{}
+
+// Compress enables connection-level compression.
+//
+// Unlike other commands, this method blocks until the command completes.
+//
+// A nil options pointer is equivalent to a zero options value.
+func (c *Client) Compress(options *CompressOptions) error {
+	upgradeDone := make(chan struct{})
+	cmd := &compressCommand{
+		upgradeDone: upgradeDone,
+	}
+	enc := c.beginCommand("COMPRESS", cmd)
+	enc.SP().Atom("DEFLATE")
+	enc.flush()
+	defer enc.end()
+
+	// The client MUST NOT send any further commands until it has seen the
+	// result of COMPRESS.
+
+	if err := cmd.Wait(); err != nil {
+		return err
+	}
+
+	// The decoder goroutine will invoke Client.upgradeCompress
+	<-upgradeDone
+	return nil
+}
+
+func (c *Client) upgradeCompress(compress *compressCommand) {
+	defer close(compress.upgradeDone)
+
+	// Drain buffered data from our bufio.Reader
+	var buf bytes.Buffer
+	if _, err := io.CopyN(&buf, c.br, int64(c.br.Buffered())); err != nil {
+		panic(err) // unreachable
+	}
+
+	conn := c.conn
+	if c.tlsConn != nil {
+		conn = c.tlsConn
+	}
+
+	var r io.Reader
+	if buf.Len() > 0 {
+		r = io.MultiReader(&buf, conn)
+	} else {
+		r = c.conn
+	}
+
+	w, err := flate.NewWriter(conn, flate.DefaultCompression)
+	if err != nil {
+		panic(err) // can only happen due to bad arguments
+	}
+
+	rw := c.options.wrapReadWriter(struct {
+		io.Reader
+		io.Writer
+	}{
+		Reader: flate.NewReader(r),
+		Writer: w,
+	})
+
+	c.br.Reset(rw)
+	// Unfortunately we can't re-use the bufio.Writer here, it races with
+	// Client.Compress
+	c.bw = bufio.NewWriter(rw)
+}
+
+type compressCommand struct {
+	cmd
+	upgradeDone chan<- struct{}
+}

imapclient/compress_test.go

@@ -0,0 +1,25 @@
+package imapclient_test
+
+import (
+	"testing"
+
+	"github.com/emersion/go-imap/v2"
+)
+
+func TestCompress(t *testing.T) {
+	client, server := newClientServerPair(t, imap.ConnStateAuthenticated)
+	defer client.Close()
+	defer server.Close()
+
+	if algos := client.Caps().CompressAlgorithms(); len(algos) == 0 {
+		t.Skipf("COMPRESS not supported")
+	}
+
+	if err := client.Compress(nil); err != nil {
+		t.Fatalf("Compress() = %v", err)
+	}
+
+	if err := client.Noop().Wait(); err != nil {
+		t.Fatalf("Noop().Wait() = %v", err)
+	}
+}

imapclient/starttls.go

@@ -57,6 +57,10 @@ func (c *Client) upgradeStartTLS(startTLS *startTLSCommand) {
 	tlsConn := tls.Client(cleartextConn, startTLS.tlsConfig)
 	rw := c.options.wrapReadWriter(tlsConn)
 
+	if c.tlsConn != nil {
+		panic("imapclient: TLS started twice")
+	}
+	c.tlsConn = tlsConn
 	c.br.Reset(rw)
 	// Unfortunately we can't re-use the bufio.Writer here, it races with
 	// Client.StartTLS

imapserver/capability.go

@@ -57,6 +57,9 @@ func (c *Conn) availableCaps() []imap.Cap {
 	} else if c.state == imap.ConnStateNotAuthenticated {
 		caps = append(caps, imap.CapLoginDisabled)
 	}
+	if c.canCompress() {
+		caps = append(caps, imap.Cap("COMPRESS=DEFLATE"))
+	}
 	if c.state == imap.ConnStateAuthenticated || c.state == imap.ConnStateSelected {
 		if available.Has(imap.CapIMAP4rev1) {
 			caps = append(caps, []imap.Cap{

imapserver/compress.go

@@ -0,0 +1,86 @@
+package imapserver
+
+import (
+	"bytes"
+	"compress/flate"
+	"io"
+
+	"github.com/emersion/go-imap/v2"
+	"github.com/emersion/go-imap/v2/internal/imapwire"
+)
+
+func (c *Conn) canCompress() bool {
+	switch c.state {
+	case imap.ConnStateAuthenticated, imap.ConnStateSelected:
+		return true // TODO
+	default:
+		return false
+	}
+}
+
+func (c *Conn) handleCompress(tag string, dec *imapwire.Decoder) error {
+	var algo string
+	if !dec.ExpectSP() || !dec.ExpectAtom(&algo) || !dec.ExpectCRLF() {
+		return dec.Err()
+	}
+
+	if !c.canCompress() {
+		return &imap.Error{
+			Type: imap.StatusResponseTypeBad,
+			Text: "COMPRESS not available",
+		}
+	}
+	if algo != "DEFLATE" {
+		return &imap.Error{
+			Type: imap.StatusResponseTypeNo,
+			Text: "Unsupported compression algorithm",
+		}
+	}
+
+	// Do not allow to write uncompressed data past this point: keep c.encMutex
+	// locked until the end
+	enc := newResponseEncoder(c)
+	defer enc.end()
+
+	err := writeStatusResp(enc.Encoder, tag, &imap.StatusResponse{
+		Type: imap.StatusResponseTypeOK,
+		Text: "Begin compression now",
+	})
+	if err != nil {
+		return err
+	}
+
+	// Drain buffered data from our bufio.Reader
+	var buf bytes.Buffer
+	if _, err := io.CopyN(&buf, c.br, int64(c.br.Buffered())); err != nil {
+		panic(err) // unreachable
+	}
+
+	var r io.Reader
+	if buf.Len() > 0 {
+		r = io.MultiReader(&buf, c.conn)
+	} else {
+		r = c.conn
+	}
+
+	c.mutex.Lock()
+	// TODO
+	c.mutex.Unlock()
+
+	w, err := flate.NewWriter(c.conn, flate.DefaultCompression)
+	if err != nil {
+		panic(err) // can only happen due to bad arguments
+	}
+
+	rw := c.server.options.wrapReadWriter(struct {
+		io.Reader
+		io.Writer
+	}{
+		Reader: flate.NewReader(r),
+		Writer: w,
+	})
+	c.br.Reset(rw)
+	c.bw.Reset(rw)
+
+	return nil
+}

imapserver/conn.go

@@ -266,6 +266,8 @@ func (c *Conn) readCommand(dec *imapwire.Decoder) error {
 		err = c.handleMove(dec, numKind)
 	case "SEARCH", "UID SEARCH":
 		err = c.handleSearch(tag, dec, numKind)
+	case "COMPRESS":
+		err = c.handleCompress(tag, dec)
 	default:
 		if c.state == imap.ConnStateNotAuthenticated {
 			// Don't allow a single unknown command before authentication to

response.go

@@ -49,6 +49,9 @@ const (
 
 	// APPENDLIMIT
 	ResponseCodeTooBig ResponseCode = "TOOBIG"
+
+	// COMPRESS
+	ResponseCodeCompressionActive ResponseCode = "COMPRESSIONACTIVE"
 )
 
 // StatusResponse is a generic status response.