imapserver: Prevent spurious unhandled panics for connections with compress=deflate that break

Writing to a connection goes through the flate library to compress. That writes
the compressed bytes to the underlying connection. But that underlying
connection is wrapped to raise a panic with an i/o error instead of returning a
normal error.  Jumping out of flate leaves the internal state of the compressor
in undefined state. So far so good. But as part of cleaning up the connection,
we could try to flush output again. Specifically: If we were writing user data,
we had switched from tracing of protocol data to tracing of user data, and we
registered a defer that restored the tracing kind and flushed (to ensure data
was traced at the right level). That flush would cause a write into the
compressor again, which could panic with an out of bounds slice access due to
its inconsistent internal state.

This fix prevents that compressor panic in two ways:

1. We wrap the flate.Writer with a moxio.FlateWriter that keeps track of
   whether a panic came out of an operation on it. If so, any further operation
   raises the same panic. This prevents access to the inconsistent internal flate
   state entirely.
2. Once we raise an i/o error, we mark the connection as broken and that makes
   flushes a no-op.

imapclient/client.go

@@ -22,8 +22,6 @@ import (
 	"reflect"
 	"strings"
 
-	"github.com/mjl-/flate"
-
 	"github.com/mjl-/mox/mlog"
 	"github.com/mjl-/mox/moxio"
 )
@@ -34,10 +32,11 @@ type Conn struct {
 	// writes through c.bw. It wraps 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
 	bw          *bufio.Writer
 	compress    bool // If compression is enabled, we must flush flateWriter and its target original bufio writer.
-	flateWriter *flate.Writer
+	flateWriter *moxio.FlateWriter
 	flateBW     *bufio.Writer
 
 	log       mlog.Log
@@ -146,11 +145,19 @@ 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
 }
 
 func (c *Conn) xflush() {
+	// Not writing any more when connection is broken.
+	if c.connBroken {
+		return
+	}
+
 	err := c.bw.Flush()
 	c.xcheckf(err, "flush")
 
@@ -173,7 +180,7 @@ func (c *Conn) Close() (rerr error) {
 	if c.conn == nil {
 		return nil
 	}
-	if c.flateWriter != nil {
+	if !c.connBroken && c.flateWriter != nil {
 		err := c.flateWriter.Close()
 		c.xcheckf(err, "close deflate writer")
 		err = c.flateBW.Flush()

imapclient/cmds.go

@@ -140,8 +140,9 @@ func (c *Conn) CompressDeflate() (untagged []Untagged, result Result, rerr error
 	c.xcheck(rerr)
 
 	c.flateBW = bufio.NewWriter(c)
-	fw, err := flate.NewWriter(c.flateBW, flate.DefaultCompression)
+	fw0, err := flate.NewWriter(c.flateBW, flate.DefaultCompression)
 	c.xcheckf(err, "deflate") // Cannot happen.
+	fw := moxio.NewFlateWriter(fw0)
 
 	c.compress = true
 	c.flateWriter = fw