fix(qr): serialize websocket I/O and keep reading past heartbeat_ack

internal/ui/login/qr/conn_test.go

@@ -0,0 +1,107 @@
+package qr
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/ayn2op/tview"
+	"github.com/gdamore/tcell/v3"
+	"github.com/gdamore/tcell/v3/vt"
+	"github.com/gorilla/websocket"
+)
+
+// gatewayModel runs Model on an open connection and quits once the login is canceled.
+type gatewayModel struct {
+	Model
+	conn *websocket.Conn
+}
+
+func (m gatewayModel) Init() tview.Cmd {
+	return func() tview.Msg { return connCreateMsg{conn: &gatewayConn{ws: m.conn}} }
+}
+
+func (m gatewayModel) Update(msg tview.Msg) (gatewayModel, tview.Cmd) {
+	var cmd tview.Cmd
+	m.Model, cmd = m.Model.Update(msg)
+	if m.status == "Login canceled on mobile" {
+		return m, tview.Quit()
+	}
+	return m, cmd
+}
+
+// serveGateway fakes the Remote Auth Gateway, heartbeating every millisecond to overlap the client's writes, and cancels after acks heartbeat acks.
+func serveGateway(t *testing.T, acks int) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer conn.Close()
+
+		if err := conn.WriteJSON(map[string]any{"op": "hello", "heartbeat_interval": 1}); err != nil {
+			return
+		}
+		shown := false
+		for {
+			var msg struct {
+				Op string `json:"op"`
+			}
+			err := conn.ReadJSON(&msg)
+			switch {
+			case err != nil:
+			case msg.Op == "heartbeat":
+				if shown {
+					acks--
+				}
+				err = conn.WriteJSON(map[string]any{"op": "heartbeat_ack"})
+			case msg.Op == "init":
+				shown = true
+				err = conn.WriteJSON(map[string]any{"op": "pending_remote_init", "fingerprint": "fingerprint"})
+			}
+			if err != nil {
+				return
+			}
+			if shown && acks <= 0 {
+				_ = conn.WriteJSON(map[string]any{"op": "cancel"})
+				return
+			}
+		}
+	}
+}
+
+func TestModelGateway(t *testing.T) {
+	t.Run("reads messages after heartbeat acks", func(t *testing.T) {
+		server := httptest.NewServer(serveGateway(t, 5))
+		t.Cleanup(server.Close)
+		conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() { conn.Close() })
+
+		screen, err := tcell.NewTerminfoScreenFromTty(vt.NewMockTerm(vt.MockOptSize{X: 80, Y: 24}))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := screen.Init(); err != nil {
+			t.Fatal(err)
+		}
+
+		done := make(chan error, 1)
+		go func() {
+			done <- tview.NewApplication(gatewayModel{Model: NewModel(), conn: conn}, tview.WithScreen(screen)).Run()
+		}()
+		select {
+		case err := <-done:
+			if err != nil {
+				t.Fatal(err)
+			}
+		case <-time.After(5 * time.Second):
+			t.Fatal("the cancel from the gateway was never read")
+		}
+	})
+}

internal/ui/login/qr/model.go

@@ -10,7 +10,6 @@ import (
 	"github.com/ayn2op/tview/richtext"
 	"github.com/ayn2op/tview/textview"
 	"github.com/gdamore/tcell/v3"
-	"github.com/gorilla/websocket"
 	"github.com/skip2/go-qrcode"
 )
 
@@ -19,7 +18,7 @@ type Model struct {
 	code        richtext.Text
 	scrollState textview.ScrollState
 
-	conn              *websocket.Conn
+	conn              *gatewayConn
 	heartbeatInterval time.Duration
 	privateKey        *rsa.PrivateKey
 	fingerprint       string
@@ -84,7 +83,7 @@ func (m *Model) update(msg tview.Msg) tview.Cmd {
 		return tview.Batch(listen(m.conn), scheduleHeartbeat(m.heartbeatInterval), generatePrivateKey())
 	case privateKeyMsg:
 		m.privateKey = msg.privateKey
-		return tview.Batch(listen(m.conn), sendInit(m.conn, m.privateKey))
+		return sendInit(m.conn, m.privateKey)
 	case nonceProofMsg:
 		return tview.Batch(listen(m.conn), sendNonceProof(m.conn, m.privateKey, msg.encryptedNonce))
 	case pendingRemoteInitMsg:
@@ -93,7 +92,7 @@ func (m *Model) update(msg tview.Msg) tview.Cmd {
 	case qrCodeMsg:
 		m.qrCode = msg.qrCode
 		m.setStatus("Scan this with the Discord mobile app to log in instantly.")
-		return listen(m.conn)
+		return nil
 	case pendingTicketMsg:
 		return tview.Batch(listen(m.conn), decryptUserPayload(m.privateKey, msg.encryptedUserPayload))
 	case userMsg:
@@ -102,10 +101,12 @@ func (m *Model) update(msg tview.Msg) tview.Cmd {
 			name += "#" + msg.discriminator
 		}
 		m.setStatus("Check your phone! Logging in as " + name)
-		return listen(m.conn)
+		return nil
 	case pendingLoginMsg:
 		m.setStatus("Authenticating...")
 		return tview.Batch(closeConn(m.conn), exchangeTicket(m.fingerprint, m.privateKey, msg.ticket))
+	case ignoredMsg:
+		return listen(m.conn)
 	case cancelMsg:
 		m.setStatus("Login canceled on mobile")
 		return closeConn(m.conn)

internal/ui/login/qr/msg.go

@@ -9,6 +9,7 @@ import (
 	json "encoding/json/v2"
 	"errors"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ayn2op/arikawa/v3/utils/httputil"
@@ -26,8 +27,20 @@ type (
 
 const remoteAuthGatewayURL = "wss://remote-auth-gateway.discord.gg/?v=2"
 
+// gatewayConn is a Remote Auth Gateway connection whose writes are serialized, since a websocket.Conn takes one writer at a time.
+type gatewayConn struct {
+	ws      *websocket.Conn
+	writeMu sync.Mutex
+}
+
+func (c *gatewayConn) writeJSON(v any) error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	return c.ws.WriteJSON(v)
+}
+
 type connCreateMsg struct {
-	conn *websocket.Conn
+	conn *gatewayConn
 }
 
 type connCloseMsg struct{}
@@ -41,14 +54,14 @@ func connect() tview.Cmd {
 		if err != nil {
 			return errMsg(err)
 		}
-		return connCreateMsg{conn: conn}
+		return connCreateMsg{conn: &gatewayConn{ws: conn}}
 	}
 }
 
-func closeConn(conn *websocket.Conn) tview.Cmd {
+func closeConn(conn *gatewayConn) tview.Cmd {
 	return func() tview.Msg {
 		if conn != nil {
-			if err := conn.Close(); err != nil {
+			if err := conn.ws.Close(); err != nil {
 				return errMsg(err)
 			}
 		}
@@ -79,13 +92,16 @@ type pendingLoginMsg struct {
 
 type cancelMsg struct{}
 
-func listen(conn *websocket.Conn) tview.Cmd {
+// ignoredMsg is a gateway message the login does not act on, such as heartbeat_ack.
+type ignoredMsg struct{}
+
+func listen(conn *gatewayConn) tview.Cmd {
 	return func() tview.Msg {
 		if conn == nil {
 			return nil
 		}
 
-		_, data, err := conn.ReadMessage()
+		_, data, err := conn.ws.ReadMessage()
 		if err != nil {
 			return errMsg(err)
 		}
@@ -122,7 +138,7 @@ func decodeMessage(data []byte) tview.Msg {
 	case "pending_login":
 		return pendingLoginMsg{ticket: payload.Ticket}
 	default:
-		return nil
+		return ignoredMsg{}
 	}
 }
 
@@ -135,7 +151,7 @@ func scheduleHeartbeat(interval time.Duration) tview.Cmd {
 	}
 }
 
-func sendHeartbeat(conn *websocket.Conn) tview.Cmd {
+func sendHeartbeat(conn *gatewayConn) tview.Cmd {
 	return func() tview.Msg {
 		if conn == nil {
 			return nil
@@ -143,7 +159,7 @@ func sendHeartbeat(conn *websocket.Conn) tview.Cmd {
 		data := struct {
 			Op string `json:"op"`
 		}{"heartbeat"}
-		if err := conn.WriteJSON(data); err != nil {
+		if err := conn.writeJSON(data); err != nil {
 			return errMsg(err)
 		}
 		return nil
@@ -164,7 +180,7 @@ func generatePrivateKey() tview.Cmd {
 	}
 }
 
-func sendInit(conn *websocket.Conn, privateKey *rsa.PrivateKey) tview.Cmd {
+func sendInit(conn *gatewayConn, privateKey *rsa.PrivateKey) tview.Cmd {
 	return func() tview.Msg {
 		if privateKey == nil {
 			return errMsg(errors.New("missing private key"))
@@ -178,14 +194,14 @@ func sendInit(conn *websocket.Conn, privateKey *rsa.PrivateKey) tview.Cmd {
 			Op               string `json:"op"`
 			EncodedPublicKey string `json:"encoded_public_key"`
 		}{"init", encodedPublicKey}
-		if err := conn.WriteJSON(data); err != nil {
+		if err := conn.writeJSON(data); err != nil {
 			return errMsg(err)
 		}
 		return nil
 	}
 }
 
-func sendNonceProof(conn *websocket.Conn, privateKey *rsa.PrivateKey, encryptedNonce string) tview.Cmd {
+func sendNonceProof(conn *gatewayConn, privateKey *rsa.PrivateKey, encryptedNonce string) tview.Cmd {
 	return func() tview.Msg {
 		decodedNonce, err := base64.StdEncoding.DecodeString(encryptedNonce)
 		if err != nil {
@@ -202,7 +218,7 @@ func sendNonceProof(conn *websocket.Conn, privateKey *rsa.PrivateKey, encryptedN
 			Op    string `json:"op"`
 			Nonce string `json:"nonce"`
 		}{"nonce_proof", encodedNonce}
-		if err := conn.WriteJSON(data); err != nil {
+		if err := conn.writeJSON(data); err != nil {
 			return errMsg(err)
 		}
 		return nil

internal/ui/login/qr/msg_test.go

@@ -14,7 +14,7 @@ func TestDecodeMessage(t *testing.T) {
 		{"pending_ticket", `{"op":"pending_ticket","encrypted_user_payload":"payload"}`, pendingTicketMsg{encryptedUserPayload: "payload"}},
 		{"cancel", `{"op":"cancel"}`, cancelMsg{}},
 		{"pending_login", `{"op":"pending_login","ticket":"ticket"}`, pendingLoginMsg{ticket: "ticket"}},
-		{"unknown", `{"op":"unknown"}`, nil},
+		{"unknown", `{"op":"unknown"}`, ignoredMsg{}},
 	}
 
 	for _, test := range tests {