fix: unblock accesspoint close during in-flight send

Author: gjermundgaraba

ap/ap.go

@@ -52,15 +52,14 @@ type Accesspoint struct {
 	encConn *shannonConn
 
 	done            chan struct{}
+	closeOnce       sync.Once
 	recvLoopOnce    sync.Once
 	recvChans       map[PacketType][]chan Packet
 	recvChansLock   sync.RWMutex
 	lastPongAck     time.Time
 	lastPongAckLock sync.Mutex
 
-	// connMu is held for writing when performing reconnection and for reading mainly when accessing welcome
-	// or sending packets. If it's not held, a valid connection (and APWelcome) is available. Be careful not to deadlock
-	// anything with this.
+	// connMu protects conn, encConn, and welcome pointer state.
 	connMu  sync.RWMutex
 	welcome *pb.APWelcome
 }
@@ -243,30 +242,34 @@ func (ap *Accesspoint) connect(ctx context.Context, creds *pb.LoginCredentials)
 }
 
 func (ap *Accesspoint) Close() {
-	ap.connMu.Lock()
-	defer ap.connMu.Unlock()
-
-	select {
-	case <-ap.done:
-		return
-	default:
-	}
-
-	close(ap.done)
-	ap.closeConnLocked()
+	ap.closeOnce.Do(func() {
+		close(ap.done)
+		ap.closeConn()
+	})
 }
 
 func (ap *Accesspoint) Send(ctx context.Context, pktType PacketType, payload []byte) error {
 	ap.connMu.RLock()
-	defer ap.connMu.RUnlock()
-
 	select {
 	case <-ap.done:
+		ap.connMu.RUnlock()
 		return ErrAccesspointClosed
 	default:
 	}
 
-	return ap.encConn.sendPacket(ctx, pktType, payload)
+	encConn := ap.encConn
+	ap.connMu.RUnlock()
+
+	if err := encConn.sendPacket(ctx, pktType, payload); err != nil {
+		select {
+		case <-ap.done:
+			return ErrAccesspointClosed
+		default:
+			return err
+		}
+	}
+
+	return nil
 }
 
 func (ap *Accesspoint) Receive(types ...PacketType) <-chan Packet {
@@ -453,14 +456,12 @@ func (ap *Accesspoint) timeSinceLastPongAck() time.Duration {
 }
 
 func (ap *Accesspoint) closeConn() {
-	ap.connMu.Lock()
-	ap.closeConnLocked()
-	ap.connMu.Unlock()
-}
+	ap.connMu.RLock()
+	conn := ap.conn
+	ap.connMu.RUnlock()
 
-func (ap *Accesspoint) closeConnLocked() {
-	if ap.conn != nil {
-		_ = ap.conn.Close()
+	if conn != nil {
+		_ = conn.Close()
 	}
 }
 

ap/ap_test.go

@@ -2,9 +2,9 @@ package ap
 
 import (
 	"context"
+	"errors"
 	"io"
 	"net"
-	"sync"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -17,36 +17,15 @@ type stubAddr string
 func (a stubAddr) Network() string { return string(a) }
 func (a stubAddr) String() string  { return string(a) }
 
-type countingConn struct {
-	mu     sync.Mutex
-	closes int
-}
-
-func (c *countingConn) Read([]byte) (int, error)         { return 0, io.EOF }
-func (c *countingConn) Write(b []byte) (int, error)      { return len(b), nil }
-func (c *countingConn) Close() error                     { c.mu.Lock(); c.closes++; c.mu.Unlock(); return nil }
-func (c *countingConn) LocalAddr() net.Addr              { return stubAddr("local") }
-func (c *countingConn) RemoteAddr() net.Addr             { return stubAddr("remote") }
-func (c *countingConn) SetDeadline(time.Time) error      { return nil }
-func (c *countingConn) SetReadDeadline(time.Time) error  { return nil }
-func (c *countingConn) SetWriteDeadline(time.Time) error { return nil }
-
-func (c *countingConn) CloseCount() int {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	return c.closes
-}
-
 type blockingConn struct {
-	started sync.Once
-	startCh chan struct{}
-	blockCh chan struct{}
+	startCh  chan struct{}
+	closedCh chan struct{}
 }
 
 func newBlockingConn() *blockingConn {
 	return &blockingConn{
-		startCh: make(chan struct{}),
-		blockCh: make(chan struct{}),
+		startCh:  make(chan struct{}),
+		closedCh: make(chan struct{}),
 	}
 }
 
@@ -57,13 +36,14 @@ func (c *blockingConn) SetDeadline(time.Time) error      { return nil }
 func (c *blockingConn) SetReadDeadline(time.Time) error  { return nil }
 func (c *blockingConn) SetWriteDeadline(time.Time) error { return nil }
 
-func (c *blockingConn) Write(b []byte) (int, error) {
-	c.started.Do(func() { close(c.startCh) })
-	<-c.blockCh
-	return len(b), nil
+func (c *blockingConn) Write([]byte) (int, error) {
+	close(c.startCh)
+	<-c.closedCh
+	return 0, io.ErrClosedPipe
 }
 
 func (c *blockingConn) Close() error {
+	close(c.closedCh)
 	return nil
 }
 
@@ -117,7 +97,7 @@ func TestDoneChannelClosesOnClose(t *testing.T) {
 	}
 }
 
-func TestCloseWaitsForInFlightSend(t *testing.T) {
+func TestCloseSignalsAndUnblocksInFlightSend(t *testing.T) {
 	conn := newBlockingConn()
 	ap := NewAccesspoint(&librespot.NullLogger{}, nil, "")
 	ap.conn = conn
@@ -141,25 +121,23 @@ func TestCloseWaitsForInFlightSend(t *testing.T) {
 	}()
 
 	select {
-	case <-closeDone:
-		t.Fatal("close returned before in-flight send finished")
-	case <-time.After(50 * time.Millisecond):
+	case <-ap.Done():
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for done channel to close")
 	}
 
-	close(conn.blockCh)
-
 	select {
-	case err := <-sendDone:
-		if err != nil {
-			t.Fatalf("send error = %v", err)
-		}
+	case <-closeDone:
 	case <-time.After(time.Second):
-		t.Fatal("timed out waiting for send to finish")
+		t.Fatal("timed out waiting for close to finish")
 	}
 
 	select {
-	case <-closeDone:
+	case err := <-sendDone:
+		if !errors.Is(err, ErrAccesspointClosed) {
+			t.Fatalf("expected ErrAccesspointClosed, got %v", err)
+		}
 	case <-time.After(time.Second):
-		t.Fatal("timed out waiting for close to finish")
+		t.Fatal("timed out waiting for send to finish")
 	}
 }

dealer/dealer.go

@@ -33,12 +33,12 @@ type Dealer struct {
 	conn *websocket.Conn
 
 	done         chan struct{}
+	closeOnce    sync.Once
 	recvLoopOnce sync.Once
 	lastPong     time.Time
 	lastPongLock sync.Mutex
 
-	// connMu is held for writing when performing reconnection and for reading when accessing the conn.
-	// If it's not held, a valid connection is available. Be careful not to deadlock anything with this.
+	// connMu protects conn pointer state.
 	connMu sync.RWMutex
 
 	messageReceivers     []messageReceiver
@@ -113,21 +113,10 @@ func (d *Dealer) connect(ctx context.Context) error {
 }
 
 func (d *Dealer) Close() {
-	d.connMu.Lock()
-	select {
-	case <-d.done:
-		d.connMu.Unlock()
-		return
-	default:
-	}
-
-	close(d.done)
-	conn := d.conn
-	d.connMu.Unlock()
-
-	if conn != nil {
-		_ = conn.Close(websocket.StatusGoingAway, "")
-	}
+	d.closeOnce.Do(func() {
+		close(d.done)
+		d.closeConn(websocket.StatusGoingAway)
+	})
 }
 
 func (d *Dealer) startReceiving() {
@@ -336,7 +325,6 @@ func (d *Dealer) closeConnRef(conn *websocket.Conn, status websocket.StatusCode)
 
 func (d *Dealer) writeConn(ctx context.Context, typ websocket.MessageType, payload []byte) (*websocket.Conn, error) {
 	d.connMu.RLock()
-
 	select {
 	case <-d.done:
 		d.connMu.RUnlock()
@@ -345,14 +333,21 @@ func (d *Dealer) writeConn(ctx context.Context, typ websocket.MessageType, paylo
 	}
 
 	conn := d.conn
+	d.connMu.RUnlock()
 
 	if conn == nil {
-		d.connMu.RUnlock()
 		return nil, fmt.Errorf("dealer connection not established")
 	}
 
 	err := conn.Write(ctx, typ, payload)
-	d.connMu.RUnlock()
+	if err != nil {
+		select {
+		case <-d.done:
+			return conn, ErrDealerClosed
+		default:
+		}
+	}
+
 	return conn, err
 }