TUN-5184: Make sure outstanding websocket write is finished, and no more writes after shutdown

websocket/connection.go

@@ -3,8 +3,10 @@ package websocket
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
+	"sync"
 	"time"
 
 	gobwas "github.com/gobwas/ws"
@@ -14,9 +16,6 @@ import (
 )
 
 const (
-	// Time allowed to write a message to the peer.
-	writeWait = 10 * time.Second
-
 	// Time allowed to read the next pong message from the peer.
 	defaultPongWait = 60 * time.Second
 
@@ -79,34 +78,20 @@ func (c *GorillaConn) SetDeadline(t time.Time) error {
 	return nil
 }
 
-// pinger simulates the websocket connection to keep it alive
-func (c *GorillaConn) pinger(ctx context.Context) {
-	ticker := time.NewTicker(defaultPingPeriod)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ticker.C:
-			if err := c.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
-				c.log.Debug().Msgf("failed to send ping message: %s", err)
-			}
-		case <-ctx.Done():
-			return
-		}
-	}
-}
-
 type Conn struct {
 	rw  io.ReadWriter
 	log *zerolog.Logger
-	// closed is a channel to indicate if Conn has been fully terminated
-	shutdownC chan struct{}
+	// writeLock makes sure
+	// 1. Only one write at a time. The pinger and Stream function can both call write.
+	// 2. Close only returns after in progress Write is finished, and no more Write will succeed after calling Close.
+	writeLock sync.Mutex
+	done      bool
 }
 
 func NewConn(ctx context.Context, rw io.ReadWriter, log *zerolog.Logger) *Conn {
 	c := &Conn{
-		rw:        rw,
-		log:       log,
-		shutdownC: make(chan struct{}),
+		rw:  rw,
+		log: log,
 	}
 	go c.pinger(ctx)
 	return c
@@ -121,16 +106,22 @@ func (c *Conn) Read(reader []byte) (int, error) {
 	return copy(reader, data), nil
 }
 
-// Write will write messages to the websocket connection
+// Write will write messages to the websocket connection.
+// It will not write to the connection after Close is called to fix TUN-5184
 func (c *Conn) Write(p []byte) (int, error) {
+	c.writeLock.Lock()
+	defer c.writeLock.Unlock()
+	if c.done {
+		return 0, errors.New("write to closed websocket connection")
+	}
 	if err := wsutil.WriteServerBinary(c.rw, p); err != nil {
 		return 0, err
 	}
+
 	return len(p), nil
 }
 
 func (c *Conn) pinger(ctx context.Context) {
-	defer close(c.shutdownC)
 	pongMessge := wsutil.Message{
 		OpCode:  gobwas.OpPong,
 		Payload: []byte{},
@@ -141,7 +132,11 @@ func (c *Conn) pinger(ctx context.Context) {
 	for {
 		select {
 		case <-ticker.C:
-			if err := wsutil.WriteServerMessage(c.rw, gobwas.OpPing, []byte{}); err != nil {
+			done, err := c.ping()
+			if done {
+				return
+			}
+			if err != nil {
 				c.log.Debug().Err(err).Msgf("failed to write ping message")
 			}
 			if err := wsutil.HandleClientControlMessage(c.rw, pongMessge); err != nil {
@@ -153,6 +148,17 @@ func (c *Conn) pinger(ctx context.Context) {
 	}
 }
 
+func (c *Conn) ping() (bool, error) {
+	c.writeLock.Lock()
+	defer c.writeLock.Unlock()
+
+	if c.done {
+		return true, nil
+	}
+
+	return false, wsutil.WriteServerMessage(c.rw, gobwas.OpPing, []byte{})
+}
+
 func (c *Conn) pingPeriod(ctx context.Context) time.Duration {
 	if val := ctx.Value(PingPeriodContextKey); val != nil {
 		if period, ok := val.(time.Duration); ok {
@@ -162,7 +168,9 @@ func (c *Conn) pingPeriod(ctx context.Context) time.Duration {
 	return defaultPingPeriod
 }
 
-// Close waits for pinger to terminate
-func (c *Conn) WaitForShutdown() {
-	<-c.shutdownC
+// Close waits for the current write to finish. Further writes will return error
+func (c *Conn) Close() {
+	c.writeLock.Lock()
+	defer c.writeLock.Unlock()
+	c.done = true
 }