imapclient: add COMPRESS

Author: emersion

imapclient/client.go

@@ -618,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)
@@ -637,6 +638,9 @@ func (c *Client) readResponse() error {
 	if startTLS != nil {
 		c.upgradeStartTLS(startTLS)
 	}
+	if compress != nil {
+		c.upgradeCompress(compress)
+	}
 
 	return nil
 }
@@ -666,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
@@ -687,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":
@@ -703,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
@@ -728,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
@@ -748,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) {
@@ -765,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)
+	}
+}

response.go

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