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