Merge support for COMPRESS extension

Author: foxcpp

README.md

@@ -135,6 +135,7 @@ includes:
 * [CHILDREN](https://tools.ietf.org/html/rfc3348)
 * [UNSELECT](https://tools.ietf.org/html/rfc3691)
 * [APPENDLIMIT](https://tools.ietf.org/html/rfc7889)
+* [COMPRESS](https://tools.ietf.org/html/rfc4978)
 
 Support for other extensions is provided via separate packages. See below.
 
@@ -146,7 +147,6 @@ Commands defined in IMAP extensions are available in other packages. See [the
 wiki](https://github.com/emersion/go-imap/wiki/Using-extensions#using-client-extensions)
 to learn how to use them.
 
-* [COMPRESS](https://github.com/emersion/go-imap-compress)
 * [ENABLE](https://github.com/emersion/go-imap-enable)
 * [ID](https://github.com/ProtonMail/go-imap-id)
 * [IDLE](https://github.com/emersion/go-imap-idle)

client/client.go

@@ -66,9 +66,10 @@ type Client struct {
 	isTLS      bool
 	serverName string
 
-	loggedOut chan struct{}
-	continues chan<- bool
-	upgrading bool
+	loggedOut    chan struct{}
+	continues    chan<- bool
+	upgrading    bool
+	isCompressed bool
 
 	handlers       []responses.Handler
 	handlersLocker sync.Mutex

client/cmd_any.go

@@ -1,16 +1,23 @@
 package client
 
 import (
+	"compress/flate"
 	"errors"
+	"net"
 
 	"github.com/emersion/go-imap"
 	"github.com/emersion/go-imap/commands"
+	"github.com/emersion/go-imap/internal"
 )
 
 // ErrAlreadyLoggedOut is returned if Logout is called when the client is
 // already logged out.
 var ErrAlreadyLoggedOut = errors.New("Already logged out")
 
+// ErrAlreadyCompress is returned by Client.Compress when compression has
+// already been enabled on the client.
+var ErrAlreadyCompressed = errors.New("COMPRESS is already enabled")
+
 // Capability requests a listing of capabilities that the server supports.
 // Capabilities are often returned by the server with the greeting or with the
 // STARTTLS and LOGIN responses, so usually explicitly requesting capabilities
@@ -86,3 +93,35 @@ func (c *Client) Logout() error {
 	}
 	return nil
 }
+
+// Compress instructs the server to use the named compression mechanism for all
+// commands and/or responses.
+func (c *Client) Compress(mech string) error {
+	if c.isCompressed {
+		return ErrAlreadyCompressed
+	}
+
+	if ok, err := c.Support("COMPRESS=" + mech); !ok || err != nil {
+		return imap.CompressUnsupportedError{Mechanism: mech}
+	}
+	if mech != imap.CompressDeflate {
+		return imap.CompressUnsupportedError{Mechanism: mech}
+	}
+
+	cmd := &commands.Compress{Mechanism: mech}
+	err := c.Upgrade(func(conn net.Conn) (net.Conn, error) {
+		if status, err := c.Execute(cmd, nil); err != nil {
+			return nil, err
+		} else if err := status.Err(); err != nil {
+			return nil, err
+		}
+
+		return internal.CreateDeflateConn(conn, flate.DefaultCompression)
+	})
+	if err != nil {
+		return err
+	}
+
+	c.isCompressed = true
+	return nil
+}

commands/compress.go

@@ -0,0 +1,33 @@
+package commands
+
+import (
+	"errors"
+
+	"github.com/emersion/go-imap"
+)
+
+// A COMPRESS command.
+type Compress struct {
+	// Name of the compression mechanism.
+	Mechanism string
+}
+
+func (cmd *Compress) Command() *imap.Command {
+	return &imap.Command{
+		Name:      "COMPRESS",
+		Arguments: []interface{}{cmd.Mechanism},
+	}
+}
+
+func (cmd *Compress) Parse(fields []interface{}) (err error) {
+	if len(fields) < 1 {
+		return errors.New("No enough arguments")
+	}
+
+	var ok bool
+	if cmd.Mechanism, ok = fields[0].(string); !ok {
+		return errors.New("Compression mechanism name must be a string")
+	}
+
+	return nil
+}

deflate.go

@@ -0,0 +1,17 @@
+package imap
+
+// A CompressUnsuppportedError is returned by Client.Compress when the provided
+// compression mechanism is not supported.
+type CompressUnsupportedError struct {
+	Mechanism string
+}
+
+func (err CompressUnsupportedError) Error() string {
+	return "COMPRESS mechanism " + err.Mechanism + " not supported"
+}
+
+// Compression algorithms for use with COMPRESS extension (RFC 4978).
+const (
+	// The DEFLATE algorithm, defined in RFC 1951.
+	CompressDeflate = "DEFLATE"
+)

internal/deflate.go

@@ -0,0 +1,62 @@
+package internal
+
+import (
+	"compress/flate"
+	"io"
+	"net"
+)
+
+type deflateConn struct {
+	net.Conn
+
+	r io.ReadCloser
+	w *flate.Writer
+}
+
+func (c *deflateConn) Read(b []byte) (int, error) {
+	return c.r.Read(b)
+}
+
+func (c *deflateConn) Write(b []byte) (int, error) {
+	return c.w.Write(b)
+}
+
+type flusher interface {
+	Flush() error
+}
+
+func (c *deflateConn) Flush() error {
+	if f, ok := c.Conn.(flusher); ok {
+		if err := f.Flush(); err != nil {
+			return err
+		}
+	}
+
+	return c.w.Flush()
+}
+
+func (c *deflateConn) Close() error {
+	if err := c.r.Close(); err != nil {
+		return err
+	}
+
+	if err := c.w.Close(); err != nil {
+		return err
+	}
+
+	return c.Conn.Close()
+}
+
+func CreateDeflateConn(c net.Conn, level int) (net.Conn, error) {
+	r := flate.NewReader(c)
+	w, err := flate.NewWriter(c, level)
+	if err != nil {
+		return nil, err
+	}
+
+	return &deflateConn{
+		Conn: c,
+		r:    r,
+		w:    w,
+	}, nil
+}

server/cmd_any.go

@@ -1,9 +1,13 @@
 package server
 
 import (
+	"compress/flate"
+	"net"
+
 	"github.com/emersion/go-imap"
 	"github.com/emersion/go-imap/backend"
 	"github.com/emersion/go-imap/commands"
+	"github.com/emersion/go-imap/internal"
 	"github.com/emersion/go-imap/responses"
 )
 
@@ -50,3 +54,25 @@ func (cmd *Logout) Handle(conn Conn) error {
 	conn.Context().State = imap.LogoutState
 	return nil
 }
+
+type Compress struct {
+	commands.Compress
+}
+
+func (cmd *Compress) Handle(conn Conn) error {
+	if cmd.Mechanism != imap.CompressDeflate {
+		return imap.CompressUnsupportedError{Mechanism: cmd.Mechanism}
+	}
+	return nil
+}
+
+func (cmd *Compress) Upgrade(conn Conn) error {
+	err := conn.Upgrade(func(conn net.Conn) (net.Conn, error) {
+		return internal.CreateDeflateConn(conn, flate.DefaultCompression)
+	})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

server/conn.go

@@ -163,7 +163,7 @@ func (c *conn) Close() error {
 }
 
 func (c *conn) Capabilities() []string {
-	caps := []string{"IMAP4rev1", "LITERAL+", "SASL-IR", "CHILDREN", "UNSELECT"}
+	caps := []string{"IMAP4rev1", "LITERAL+", "SASL-IR", "CHILDREN", "UNSELECT", "COMPRESS=DEFLATE"}
 
 	appendLimitSet := false
 	if c.ctx.State == imap.AuthenticatedState {

server/server.go

@@ -187,6 +187,7 @@ func New(bkd backend.Backend) *Server {
 		"UID":     func() Handler { return &Uid{} },
 
 		"UNSELECT": func() Handler { return &Unselect{} },
+		"COMPRESS": func() Handler { return &Compress{} },
 	}
 
 	return s