@yume-chan/adb 0.0.21 → 0.0.23

package/src/daemon/dispatcher.ts

@@ -1,4 +1,8 @@
-import { AsyncOperationManager, PromiseResolver } from "@yume-chan/async";
+import {
+    AsyncOperationManager,
+    PromiseResolver,
+    delay,
+} from "@yume-chan/async";
 import type {
     Consumable,
     ReadableWritablePair,
@@ -9,10 +13,10 @@ import {
     ConsumableWritableStream,
     WritableStream,
 } from "@yume-chan/stream-extra";
-import { EMPTY_UINT8_ARRAY } from "@yume-chan/struct";
+import { EMPTY_UINT8_ARRAY, NumberFieldType } from "@yume-chan/struct";
 
 import type { AdbIncomingSocketHandler, AdbSocket, Closeable } from "../adb.js";
-import { decodeUtf8, encodeUtf8, unreachable } from "../utils/index.js";
+import { decodeUtf8, encodeUtf8 } from "../utils/index.js";
 
 import type { AdbPacketData, AdbPacketInit } from "./packet.js";
 import { AdbCommand, calculateChecksum } from "./packet.js";
@@ -28,10 +32,23 @@ export interface AdbPacketDispatcherOptions {
      */
     appendNullToServiceString: boolean;
     maxPayloadSize: number;
+    /**
+     * The number of bytes the device can send before receiving an ack packet.
+
+     * Set to 0 or any negative value to disable delayed ack.
+     * Otherwise the value must be in the range of unsigned 32-bit integer.
+     */
+    initialDelayedAckBytes: number;
     /**
      * Whether to preserve the connection open after the `AdbPacketDispatcher` is closed.
      */
     preserveConnection?: boolean | undefined;
+    debugSlowRead?: boolean | undefined;
+}
+
+interface SocketOpenResult {
+    remoteId: number;
+    availableWriteBytes: number;
 }
 
 /**
@@ -74,36 +91,28 @@ export class AdbPacketDispatcher implements Closeable {
         options: AdbPacketDispatcherOptions,
     ) {
         this.options = options;
+        // Don't allow negative values in dispatcher
+        if (this.options.initialDelayedAckBytes < 0) {
+            this.options.initialDelayedAckBytes = 0;
+        }
 
         connection.readable
             .pipeTo(
                 new WritableStream({
                     write: async (packet) => {
                         switch (packet.command) {
-                            case AdbCommand.OK:
-                                this.#handleOk(packet);
-                                break;
                             case AdbCommand.Close:
                                 await this.#handleClose(packet);
                                 break;
-                            case AdbCommand.Write:
-                                if (this.#sockets.has(packet.arg1)) {
-                                    await this.#sockets
-                                        .get(packet.arg1)!
-                                        .enqueue(packet.payload);
-                                    await this.sendPacket(
-                                        AdbCommand.OK,
-                                        packet.arg1,
-                                        packet.arg0,
-                                    );
-                                    break;
-                                }
-                                throw new Error(
-                                    `Unknown local socket id: ${packet.arg1}`,
-                                );
+                            case AdbCommand.Okay:
+                                this.#handleOkay(packet);
+                                break;
                             case AdbCommand.Open:
                                 await this.#handleOpen(packet);
                                 break;
+                            case AdbCommand.Write:
+                                await this.#handleWrite(packet);
+                                break;
                             default:
                                 // Junk data may only appear in the authentication phase,
                                 // since the dispatcher only works after authentication,
@@ -137,24 +146,6 @@ export class AdbPacketDispatcher implements Closeable {
         this.#writer = connection.writable.getWriter();
     }
 
-    #handleOk(packet: AdbPacketData) {
-        if (this.#initializers.resolve(packet.arg1, packet.arg0)) {
-            // Device successfully created the socket
-            return;
-        }
-
-        const socket = this.#sockets.get(packet.arg1);
-        if (socket) {
-            // Device has received last `WRTE` to the socket
-            socket.ack();
-            return;
-        }
-
-        // Maybe the device is responding to a packet of last connection
-        // Tell the device to close the socket
-        void this.sendPacket(AdbCommand.Close, packet.arg1, packet.arg0);
-    }
-
     async #handleClose(packet: AdbPacketData) {
         // If the socket is still pending
         if (
@@ -182,15 +173,8 @@ export class AdbPacketDispatcher implements Closeable {
         // Ignore `arg0` and search for the socket
         const socket = this.#sockets.get(packet.arg1);
         if (socket) {
-            // The device want to close the socket
-            if (!socket.closed) {
-                await this.sendPacket(
-                    AdbCommand.Close,
-                    packet.arg1,
-                    packet.arg0,
-                );
-            }
-            await socket.dispose();
+            await socket.close();
+            socket.dispose();
             this.#sockets.delete(packet.arg1);
             return;
         }
@@ -200,16 +184,59 @@ export class AdbPacketDispatcher implements Closeable {
         // the device may also respond with two `CLSE` packets.
     }
 
-    addReverseTunnel(service: string, handler: AdbIncomingSocketHandler) {
-        this.#incomingSocketHandlers.set(service, handler);
-    }
+    #handleOkay(packet: AdbPacketData) {
+        let ackBytes: number;
+        if (this.options.initialDelayedAckBytes !== 0) {
+            if (packet.payload.byteLength !== 4) {
+                throw new Error(
+                    "Invalid OKAY packet. Payload size should be 4",
+                );
+            }
+            ackBytes = NumberFieldType.Uint32.deserialize(packet.payload, true);
+        } else {
+            if (packet.payload.byteLength !== 0) {
+                throw new Error(
+                    "Invalid OKAY packet. Payload size should be 0",
+                );
+            }
+            ackBytes = Infinity;
+        }
 
-    removeReverseTunnel(address: string) {
-        this.#incomingSocketHandlers.delete(address);
+        if (
+            this.#initializers.resolve(packet.arg1, {
+                remoteId: packet.arg0,
+                availableWriteBytes: ackBytes,
+            } satisfies SocketOpenResult)
+        ) {
+            // Device successfully created the socket
+            return;
+        }
+
+        const socket = this.#sockets.get(packet.arg1);
+        if (socket) {
+            // When delayed ack is enabled, `ackBytes` is a positive number represents
+            // how many bytes the device has received from this socket.
+            // When delayed ack is disabled, `ackBytes` is always `Infinity` represents
+            // the device has received last `WRTE` packet from the socket.
+            socket.ack(ackBytes);
+            return;
+        }
+
+        // Maybe the device is responding to a packet of last connection
+        // Tell the device to close the socket
+        void this.sendPacket(AdbCommand.Close, packet.arg1, packet.arg0);
     }
 
-    clearReverseTunnels() {
-        this.#incomingSocketHandlers.clear();
+    #sendOkay(localId: number, remoteId: number, ackBytes: number) {
+        let payload: Uint8Array;
+        if (this.options.initialDelayedAckBytes !== 0) {
+            payload = new Uint8Array(4);
+            new DataView(payload.buffer).setUint32(0, ackBytes, true);
+        } else {
+            payload = EMPTY_UINT8_ARRAY;
+        }
+
+        return this.sendPacket(AdbCommand.Okay, localId, remoteId, payload);
     }
 
     async #handleOpen(packet: AdbPacketData) {
@@ -219,9 +246,20 @@ export class AdbPacketDispatcher implements Closeable {
         this.#initializers.resolve(localId, undefined);
 
         const remoteId = packet.arg0;
-        let service = decodeUtf8(packet.payload);
-        if (service.endsWith("\0")) {
-            service = service.substring(0, service.length - 1);
+        let initialDelayedAckBytes = packet.arg1;
+        const service = decodeUtf8(packet.payload);
+
+        if (this.options.initialDelayedAckBytes === 0) {
+            if (initialDelayedAckBytes !== 0) {
+                throw new Error("Invalid OPEN packet. arg1 should be 0");
+            }
+            initialDelayedAckBytes = Infinity;
+        } else {
+            if (initialDelayedAckBytes === 0) {
+                throw new Error(
+                    "Invalid OPEN packet. arg1 should be greater than 0",
+                );
+            }
         }
 
         const handler = this.#incomingSocketHandlers.get(service);
@@ -237,26 +275,66 @@ export class AdbPacketDispatcher implements Closeable {
             localCreated: false,
             service,
         });
+        controller.ack(initialDelayedAckBytes);
 
         try {
             await handler(controller.socket);
             this.#sockets.set(localId, controller);
-            await this.sendPacket(AdbCommand.OK, localId, remoteId);
+            await this.#sendOkay(
+                localId,
+                remoteId,
+                this.options.initialDelayedAckBytes,
+            );
         } catch (e) {
             await this.sendPacket(AdbCommand.Close, 0, remoteId);
         }
     }
 
+    async #handleWrite(packet: AdbPacketData) {
+        const socket = this.#sockets.get(packet.arg1);
+        if (!socket) {
+            throw new Error(`Unknown local socket id: ${packet.arg1}`);
+        }
+
+        let handled = false;
+        await Promise.race([
+            delay(5000).then(() => {
+                if (this.options.debugSlowRead && !handled) {
+                    throw new Error(
+                        `packet for \`${socket.service}\` not handled in 5 seconds`,
+                    );
+                }
+            }),
+            (async () => {
+                await socket.enqueue(packet.payload);
+                await this.#sendOkay(
+                    packet.arg1,
+                    packet.arg0,
+                    packet.payload.length,
+                );
+                handled = true;
+            })(),
+        ]);
+
+        return;
+    }
+
     async createSocket(service: string): Promise<AdbSocket> {
         if (this.options.appendNullToServiceString) {
             service += "\0";
         }
 
-        const [localId, initializer] = this.#initializers.add<number>();
-        await this.sendPacket(AdbCommand.Open, localId, 0, service);
+        const [localId, initializer] =
+            this.#initializers.add<SocketOpenResult>();
+        await this.sendPacket(
+            AdbCommand.Open,
+            localId,
+            this.options.initialDelayedAckBytes,
+            service,
+        );
 
-        // Fulfilled by `handleOk`
-        const remoteId = await initializer;
+        // Fulfilled by `handleOkay`
+        const { remoteId, availableWriteBytes } = await initializer;
         const controller = new AdbDaemonSocketController({
             dispatcher: this,
             localId,
@@ -264,11 +342,24 @@ export class AdbPacketDispatcher implements Closeable {
             localCreated: true,
             service,
         });
+        controller.ack(availableWriteBytes);
         this.#sockets.set(localId, controller);
 
         return controller.socket;
     }
 
+    addReverseTunnel(service: string, handler: AdbIncomingSocketHandler) {
+        this.#incomingSocketHandlers.set(service, handler);
+    }
+
+    removeReverseTunnel(address: string) {
+        this.#incomingSocketHandlers.delete(address);
+    }
+
+    clearReverseTunnels() {
+        this.#incomingSocketHandlers.clear();
+    }
+
     async sendPacket(
         command: AdbCommand,
         arg0: number,
@@ -307,7 +398,7 @@ export class AdbPacketDispatcher implements Closeable {
         this.#closed = true;
 
         this.#readAbortController.abort();
-        if (this.options.preserveConnection ?? false) {
+        if (this.options.preserveConnection) {
             this.#writer.releaseLock();
         } else {
             await this.#writer.close();
@@ -318,7 +409,7 @@ export class AdbPacketDispatcher implements Closeable {
 
     #dispose() {
         for (const socket of this.#sockets.values()) {
-            socket.dispose().catch(unreachable);
+            socket.dispose();
         }
 
         this.#disconnected.resolve();

package/src/daemon/packet.ts

@@ -5,7 +5,7 @@ export enum AdbCommand {
     Auth = 0x48545541, // 'AUTH'
     Close = 0x45534c43, // 'CLSE'
     Connect = 0x4e584e43, // 'CNXN'
-    OK = 0x59414b4f, // 'OKAY'
+    Okay = 0x59414b4f, // 'OKAY'
     Open = 0x4e45504f, // 'OPEN'
     Write = 0x45545257, // 'WRTE'
 }

package/src/daemon/socket.ts

@@ -1,21 +1,19 @@
 import { PromiseResolver } from "@yume-chan/async";
 import type { Disposable } from "@yume-chan/event";
 import type {
+    AbortSignal,
     Consumable,
     PushReadableStreamController,
     ReadableStream,
-    ReadableWritablePair,
     WritableStream,
+    WritableStreamDefaultController,
 } from "@yume-chan/stream-extra";
 import {
     ConsumableWritableStream,
-    DistributionStream,
-    DuplexStreamFactory,
     PushReadableStream,
-    pipeFrom,
 } from "@yume-chan/stream-extra";
 
-import type { Closeable } from "../adb.js";
+import type { AdbSocket } from "../adb.js";
 
 import type { AdbPacketDispatcher } from "./dispatcher.js";
 import { AdbCommand } from "./packet.js";
@@ -36,11 +34,7 @@ export interface AdbDaemonSocketConstructionOptions
 }
 
 export class AdbDaemonSocketController
-    implements
-        AdbDaemonSocketInfo,
-        ReadableWritablePair<Uint8Array, Consumable<Uint8Array>>,
-        Closeable,
-        Disposable
+    implements AdbDaemonSocketInfo, AdbSocket, Disposable
 {
     readonly #dispatcher!: AdbPacketDispatcher;
 
@@ -49,25 +43,20 @@ export class AdbDaemonSocketController
     readonly localCreated!: boolean;
     readonly service!: string;
 
-    #duplex: DuplexStreamFactory<Uint8Array, Consumable<Uint8Array>>;
-
     #readable: ReadableStream<Uint8Array>;
     #readableController!: PushReadableStreamController<Uint8Array>;
     get readable() {
         return this.#readable;
     }
 
-    #writePromise: PromiseResolver<void> | undefined;
+    #writableController!: WritableStreamDefaultController;
     readonly writable: WritableStream<Consumable<Uint8Array>>;
 
     #closed = false;
-    /**
-     * Whether the socket is half-closed (i.e. the local side initiated the close).
-     *
-     * It's only used by dispatcher to avoid sending another `CLSE` packet to remote.
-     */
+
+    #closedPromise = new PromiseResolver<void>();
     get closed() {
-        return this.#closed;
+        return this.#closedPromise.promise;
     }
 
     #socket: AdbDaemonSocket;
@@ -75,6 +64,17 @@ export class AdbDaemonSocketController
         return this.#socket;
     }
 
+    #availableWriteBytesChanged: PromiseResolver<void> | undefined;
+    /**
+     * When delayed ack is disabled, can be `Infinity` if the socket is ready to write.
+     * Exactly one packet can be written no matter how large it is. Or `-1` if the socket
+     * is waiting for ack.
+     *
+     * When delayed ack is enabled, a non-negative finite number indicates the number of
+     * bytes that can be written to the socket before receiving an ack.
+     */
+    #availableWriteBytes = 0;
+
     constructor(options: AdbDaemonSocketConstructionOptions) {
         this.#dispatcher = options.dispatcher;
         this.localId = options.localId;
@@ -82,88 +82,110 @@ export class AdbDaemonSocketController
         this.localCreated = options.localCreated;
         this.service = options.service;
 
-        // Check this image to help you understand the stream graph
-        // cspell: disable-next-line
-        // https://www.plantuml.com/plantuml/png/TL0zoeGm4ErpYc3l5JxyS0yWM6mX5j4C6p4cxcJ25ejttuGX88ZftizxUKmJI275pGhXl0PP_UkfK_CAz5Z2hcWsW9Ny2fdU4C1f5aSchFVxA8vJjlTPRhqZzDQMRB7AklwJ0xXtX0ZSKH1h24ghoKAdGY23FhxC4nS2pDvxzIvxb-8THU0XlEQJ-ZB7SnXTAvc_LhOckhMdLBnbtndpb-SB7a8q2SRD_W00
-
-        this.#duplex = new DuplexStreamFactory<
-            Uint8Array,
-            Consumable<Uint8Array>
-        >({
-            close: async () => {
-                this.#closed = true;
-
-                await this.#dispatcher.sendPacket(
-                    AdbCommand.Close,
-                    this.localId,
-                    this.remoteId
-                );
-
-                // Don't `dispose` here, we need to wait for `CLSE` response packet.
-                return false;
+        this.#readable = new PushReadableStream((controller) => {
+            this.#readableController = controller;
+        });
+
+        this.writable = new ConsumableWritableStream<Uint8Array>({
+            start: (controller) => {
+                this.#writableController = controller;
             },
-            dispose: () => {
-                // Error out the pending writes
-                this.#writePromise?.reject(new Error("Socket closed"));
+            write: async (data, controller) => {
+                const size = data.length;
+                const chunkSize = this.#dispatcher.options.maxPayloadSize;
+                for (
+                    let start = 0, end = chunkSize;
+                    start < size;
+                    start = end, end += chunkSize
+                ) {
+                    const chunk = data.subarray(start, end);
+                    await this.#writeChunk(chunk, controller.signal);
+                }
             },
         });
 
-        this.#readable = this.#duplex.wrapReadable(
-            new PushReadableStream(
-                (controller) => {
-                    this.#readableController = controller;
-                },
-                {
-                    highWaterMark: options.highWaterMark ?? 16 * 1024,
-                    size(chunk) {
-                        return chunk.byteLength;
-                    },
-                }
-            )
-        );
+        this.#socket = new AdbDaemonSocket(this);
+    }
 
-        this.writable = pipeFrom(
-            this.#duplex.createWritable(
-                new ConsumableWritableStream<Uint8Array>({
-                    write: async (chunk) => {
-                        // Wait for an ack packet
-                        this.#writePromise = new PromiseResolver();
-                        await this.#dispatcher.sendPacket(
-                            AdbCommand.Write,
-                            this.localId,
-                            this.remoteId,
-                            chunk
-                        );
-                        await this.#writePromise.promise;
-                    },
-                })
-            ),
-            new DistributionStream(this.#dispatcher.options.maxPayloadSize)
-        );
+    async #writeChunk(data: Uint8Array, signal: AbortSignal) {
+        const length = data.byteLength;
+        while (this.#availableWriteBytes < length) {
+            // Only one lock is required because Web Streams API guarantees
+            // that `write` is not reentrant.
+            const resolver = new PromiseResolver<void>();
+            signal.addEventListener("abort", () => {
+                resolver.reject(signal.reason);
+            });
 
-        this.#socket = new AdbDaemonSocket(this);
+            this.#availableWriteBytesChanged = resolver;
+            await resolver.promise;
+        }
+
+        if (this.#availableWriteBytes === Infinity) {
+            this.#availableWriteBytes = -1;
+        } else {
+            this.#availableWriteBytes -= length;
+        }
+
+        await this.#dispatcher.sendPacket(
+            AdbCommand.Write,
+            this.localId,
+            this.remoteId,
+            data,
+        );
     }
 
     async enqueue(data: Uint8Array) {
-        // Consumer may abort the `ReadableStream` to close the socket,
-        // it's OK to throw away further packets in this case.
+        // Consumers can `cancel` the `readable` if they are not interested in future data.
+        // Throw away the data if that happens.
         if (this.#readableController.abortSignal.aborted) {
             return;
         }
 
-        await this.#readableController.enqueue(data);
+        try {
+            await this.#readableController.enqueue(data);
+        } catch (e) {
+            if (this.#readableController.abortSignal.aborted) {
+                return;
+            }
+            throw e;
+        }
     }
 
-    ack() {
-        this.#writePromise?.resolve();
+    public ack(bytes: number) {
+        this.#availableWriteBytes += bytes;
+        this.#availableWriteBytesChanged?.resolve();
     }
 
     async close(): Promise<void> {
-        await this.#duplex.close();
+        if (this.#closed) {
+            return;
+        }
+        this.#closed = true;
+
+        this.#availableWriteBytesChanged?.reject(new Error("Socket closed"));
+
+        try {
+            this.#writableController.error(new Error("Socket closed"));
+        } catch {
+            // ignore
+        }
+
+        await this.#dispatcher.sendPacket(
+            AdbCommand.Close,
+            this.localId,
+            this.remoteId,
+        );
     }
 
     dispose() {
-        return this.#duplex.dispose();
+        try {
+            this.#readableController.close();
+        } catch {
+            // ignore
+        }
+
+        this.#closedPromise.resolve();
     }
 }
 
@@ -175,11 +197,7 @@ export class AdbDaemonSocketController
  * `socket.writable.abort()`, `socket.writable.getWriter().abort()`,
  * `socket.writable.close()` or `socket.writable.getWriter().close()`.
  */
-export class AdbDaemonSocket
-    implements
-        AdbDaemonSocketInfo,
-        ReadableWritablePair<Uint8Array, Consumable<Uint8Array>>
-{
+export class AdbDaemonSocket implements AdbDaemonSocketInfo, AdbSocket {
     #controller: AdbDaemonSocketController;
 
     get localId(): number {
@@ -202,7 +220,7 @@ export class AdbDaemonSocket
         return this.#controller.writable;
     }
 
-    get closed(): boolean {
+    get closed(): Promise<void> {
         return this.#controller.closed;
     }