fixup: sessions client hanshake abort [discuss]

trezor · Apr 16, 2023 · de1454b · de1454b
1 parent 54325ac
commit de1454b
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 13 deletions.
diff --git a/packages/transport/src/sessions/client.ts b/packages/transport/src/sessions/client.ts
@@ -11,8 +11,14 @@ import {
 } from './types';
 import { SessionsBackground } from './background';
 import * as ERRORS from '../errors';
+import { scheduleAction } from '../utils/scheduleAction';
 
-type HandshakeTimeoutResponse = ErrorGeneric<typeof ERRORS.SESSION_BACKGROUND_TIMEOUT>;
+type HandshakeTimeoutResponse = ErrorGeneric<
+    | typeof ERRORS.SESSION_BACKGROUND_TIMEOUT
+    | typeof ERRORS.ABORTED_BY_SIGNAL
+    | typeof ERRORS.ABORTED_BY_TIMEOUT
+    | typeof ERRORS.UNEXPECTED_ERROR
+>;
 
 type RegisterCallbackOnDescriptorsChange = (callback: (descriptor: Descriptor[]) => any) => void;
 
@@ -56,19 +62,35 @@ export class SessionsClient extends TypedEmitter<{
         }
     }
 
-    async handshake() {
+    async handshake({ signal }: { signal: AbortSignal }) {
+        let hanshakeTimeout: ReturnType<typeof setTimeout> | undefined;
+        const hanshakeTimeoutPromise = scheduleAction<HandshakeTimeoutResponse>(
+            () => {
+                return new Promise(resolve => {
+                    hanshakeTimeout = setTimeout(
+                        () =>
+                            resolve({
+                                success: false,
+                                error: ERRORS.SESSION_BACKGROUND_TIMEOUT,
+                            }),
+                        5000,
+                    );
+                });
+            },
+            {
+                signal,
+            },
+        );
+
+        signal.addEventListener('abort', () => {
+            if (hanshakeTimeout) {
+                clearTimeout(hanshakeTimeout);
+            }
+        });
+
         const response = await Promise.race<HandshakeResponse | HandshakeTimeoutResponse>([
             this.request({ type: 'handshake', payload: undefined }),
-            new Promise(resolve => {
-                setTimeout(
-                    () =>
-                        resolve({
-                            success: false,
-                            error: ERRORS.SESSION_BACKGROUND_TIMEOUT,
-                        }),
-                    5000,
-                );
-            }),
+            hanshakeTimeoutPromise,
         ]);
         return response;
     }

diff --git a/packages/transport/src/transports/abstractUsb.ts b/packages/transport/src/transports/abstractUsb.ts
@@ -33,7 +33,7 @@ export abstract class AbstractUsbTransport extends AbstractTransport {
 
         return scheduleAction(
             async () => {
-                const handshakeRes = await this.sessionsClient.handshake();
+                const handshakeRes = await this.sessionsClient.handshake({  signal: localAbortController.signal });
 
                 if (!handshakeRes.success) {
                     return this.error({ error: handshakeRes.error });