Gateway: Fixed the double Close and Gateway ReconnectOP bugs

utils/wsutil/conn.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"io"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/diamondburned/arikawa/utils/json"
@@ -18,6 +19,9 @@ const CopyBufferSize = 2048
 // CloseDeadline controls the deadline to wait for sending the Close frame.
 var CloseDeadline = time.Second
 
+// ErrWebsocketClosed is returned if the websocket is already closed.
+var ErrWebsocketClosed = errors.New("Websocket is closed.")
+
 // Connection is an interface that abstracts around a generic Websocket driver.
 // This connection expects the driver to handle compression by itself, including
 // modifying the connection URL.
@@ -45,7 +49,6 @@ type Conn struct {
 	json.Driver
 
 	dialer *websocket.Dialer
-	// mut    sync.RWMutex
 	events chan Event
 
 	// write channels
@@ -53,6 +56,10 @@ type Conn struct {
 	errors chan error
 
 	buf  bytes.Buffer
+	zlib io.ReadCloser // (compress/zlib).reader
+
+	// nil until Dial().
+	closeOnce *sync.Once
 
 	// zlib *zlib.Inflator // zlib.NewReader
 	// buf  []byte         // io.Copy buffer
@@ -91,6 +98,9 @@ func (c *Conn) Dial(ctx context.Context, addr string) error {
 		return errors.Wrap(err, "Failed to dial WS")
 	}
 
+	// Set up the closer.
+	c.closeOnce = &sync.Once{}
+
 	c.events = make(chan Event)
 	go c.readLoop()
 
@@ -173,13 +183,21 @@ func (c *Conn) handle() ([]byte, error) {
 
 	if t == websocket.BinaryMessage {
 		// Probably a zlib payload
+
+		if c.zlib == nil {
 			z, err := zlib.NewReader(r)
 			if err != nil {
 				return nil, errors.Wrap(err, "Failed to create a zlib reader")
 			}
+			c.zlib = z
+		} else {
+			if err := c.zlib.(zlib.Resetter).Reset(r, nil); err != nil {
+				return nil, errors.Wrap(err, "Failed to reset zlib reader")
+			}
+		}
 
-		defer z.Close()
+		defer c.zlib.Close()
-		r = z
+		r = c.zlib
 	}
 
 	return readAll(&c.buf, r)
@@ -209,27 +227,27 @@ func (c *Conn) handle() ([]byte, error) {
 }
 
 func (c *Conn) Send(b []byte) error {
-	// Don't send a nil byte slice. That would confuse the write loop.
-	if b == nil {
-		return nil
-	}
-
 	// If websocket is already closed.
 	if c.writes == nil {
-		return errors.New("Websocket is closed.")
+		return ErrWebsocketClosed
 	}
 
 	c.writes <- b
 	return <-c.errors
 }
 
-func (c *Conn) Close() error {
+func (c *Conn) Close() (err error) {
+	// Use a sync.Once to guarantee that other Close() calls block until the
+	// main call is done. It also prevents future calls.
+	c.closeOnce.Do(func() {
 		// Close c.writes. This should trigger the websocket to close itself.
 		close(c.writes)
+		// Mark c.writes as empty.
+		c.writes = nil
 
 		// Wait for the write loop to exit by flusing the errors channel.
-	var err = <-c.errors
+		err = <-c.errors     // get close error
-	for range c.errors {
+		for range c.errors { // then flush
 		}
 
 		// Flush all events before closing the channel. This will return as soon as
@@ -242,6 +260,7 @@ func (c *Conn) Close() error {
 
 		// Mark c.Conn as empty.
 		c.Conn = nil
+	})
 
 	return err
 }