Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions uplink/websocket.go
Original file line number Diff line number Diff line change
Expand Up @@ -105,11 +105,14 @@ func (c *WebSocketConn) StartPingLoop(ctx context.Context) {
}
pctx, cancel := context.WithCancel(ctx)
c.pingCancel = cancel
c.pingDone = make(chan struct{})
done := make(chan struct{})
c.pingDone = done
c.pingMu.Unlock()

go func() {
defer close(c.pingDone)
// Close the local done channel, not c.pingDone: stopPing nils the field
// before waiting, so defer close(c.pingDone) would panic on a nil channel.
defer close(done)
t := time.NewTicker(c.pingInterval)
defer t.Stop()
for {
Expand Down
81 changes: 81 additions & 0 deletions uplink/websocket_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
package uplink

import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"

"github.com/gorilla/websocket"
)

func TestWebSocketConn_StartPingLoopCloseNoPanic(t *testing.T) {
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
defer ws.Close()
conn := NewWebSocketConn(ws, DefaultMaxMessageSize)
conn.pingInterval = 10 * time.Millisecond
conn.StartPingLoop(r.Context())
// Keep the connection open briefly so the ping loop is running,
// then Close (stopPing) races with the goroutine exit.
time.Sleep(30 * time.Millisecond)
_ = conn.Close()
}))
defer srv.Close()

wsURL := "ws" + srv.URL[len("http"):]
cli, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatal(err)
}
defer cli.Close()

// Drain until server closes; panic in the ping loop would fail the test process.
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
_ = cli.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
if _, _, err := cli.ReadMessage(); err != nil {
return
}
}
}

func TestWebSocketConn_stopPingIdempotent(t *testing.T) {
upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
done := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer close(done)
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
t.Error(err)
return
}
conn := NewWebSocketConn(ws, DefaultMaxMessageSize)
conn.pingInterval = time.Hour
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
conn.StartPingLoop(ctx)
conn.stopPing()
conn.stopPing() // second stop must not hang or panic
_ = conn.ws.Close()
}))
defer srv.Close()

wsURL := "ws" + srv.URL[len("http"):]
cli, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatal(err)
}
cli.Close()

select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("server handler did not finish")
}
}