diff --git a/internal/ui/login/qr/conn_test.go b/internal/ui/login/qr/conn_test.go new file mode 100644 index 0000000..2c9d9b0 --- /dev/null +++ b/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") + } + }) +} diff --git a/internal/ui/login/qr/model.go b/internal/ui/login/qr/model.go index 851ec56..779895f 100644 --- a/internal/ui/login/qr/model.go +++ b/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) diff --git a/internal/ui/login/qr/msg.go b/internal/ui/login/qr/msg.go index 21192bc..7a696b1 100644 --- a/internal/ui/login/qr/msg.go +++ b/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 diff --git a/internal/ui/login/qr/msg_test.go b/internal/ui/login/qr/msg_test.go index 8d27156..5caf5cf 100644 --- a/internal/ui/login/qr/msg_test.go +++ b/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 {