From 00bb0063984c32414ffd9e0117a21ea8f3eefeef Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Sat, 8 Jun 2019 20:28:37 +0300
Subject: [PATCH] client: cancel pending literal writes after status response

Fixes #113
---
 client/client.go | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/client/client.go b/client/client.go
index 16a2ee80..9699652c 100644
--- a/client/client.go
+++ b/client/client.go
@@ -68,6 +68,7 @@ type Client struct {
 	serverName string
 
 	loggedOut chan struct{}
+	continues chan<- bool
 	upgrading bool
 
 	handlers       []responses.Handler
@@ -231,6 +232,11 @@ func (c *Client) execute(cmdr imap.Commander, h responses.Handler) (*imap.Status
 				// Wait for upgrade to finish.
 				c.conn.Wait()
 			}
+			// Cancel any pending literal write
+			select {
+			case c.continues <- false:
+			default:
+			}
 			return errUnregisterHandler
 		}
 
@@ -322,11 +328,11 @@ func (c *Client) Execute(cmdr imap.Commander, h responses.Handler) (*imap.Status
 	return c.execute(cmdr, h)
 }
 
-func (c *Client) handleContinuationReqs(continues chan<- bool) {
+func (c *Client) handleContinuationReqs() {
 	c.registerHandler(responses.HandlerFunc(func(resp imap.Resp) error {
 		if _, ok := resp.(*imap.ContinuationReq); ok {
 			go func() {
-				continues <- true
+				c.continues <- true
 			}()
 			return nil
 		}
@@ -572,11 +578,12 @@ func New(conn net.Conn) (*Client, error) {
 	c := &Client{
 		conn:      imap.NewConn(conn, r, w),
 		loggedOut: make(chan struct{}),
+		continues: continues,
 		state:     imap.ConnectingState,
 		ErrorLog:  log.New(os.Stderr, "imap/client: ", log.LstdFlags),
 	}
 
-	c.handleContinuationReqs(continues)
+	c.handleContinuationReqs()
 	c.handleUnilateral()
 	err := c.handleGreetAndStartReading()
 	return c, err