feat(adb): change how to close a socket

This commit is contained in:
Simon Chan 2023-10-16 13:23:59 +08:00
parent 1aa7a92d2c
commit e45fb2ed55
No known key found for this signature in database
GPG key ID: A8B69F750B9BCEDD
6 changed files with 157 additions and 132 deletions

View file

@ -71,6 +71,10 @@ export class AdbSyncSocketLocked implements AsyncExactReadable {
this.#combiner.flush();
this.#socketLock.notifyOne();
}
async close() {
await this.#readable.cancel();
}
}
export class AdbSyncSocket {
@ -94,6 +98,7 @@ export class AdbSyncSocket {
}
async close() {
await this.#locked.close();
await this.#socket.close();
}
}

View file

@ -1,4 +1,8 @@
import { AsyncOperationManager, PromiseResolver } from "@yume-chan/async";
import {
AsyncOperationManager,
PromiseResolver,
delay,
} from "@yume-chan/async";
import type {
Consumable,
ReadableWritablePair,
@ -12,7 +16,7 @@ import {
import { EMPTY_UINT8_ARRAY } 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";
@ -80,18 +84,18 @@ export class AdbPacketDispatcher implements Closeable {
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:
await this.#handleWrite(packet);
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,
@ -125,24 +129,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 (
@ -170,15 +156,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;
}
@ -188,27 +167,22 @@ export class AdbPacketDispatcher implements Closeable {
// the device may also respond with two `CLSE` packets.
}
async #handleWrite(packet: AdbPacketData) {
const socket = this.#sockets.get(packet.arg1);
if (!socket) {
throw new Error(`Unknown local socket id: ${packet.arg1}`);
}
await socket.enqueue(packet.payload);
await this.sendPacket(AdbCommand.OK, packet.arg1, packet.arg0);
#handleOkay(packet: AdbPacketData) {
if (this.#initializers.resolve(packet.arg1, packet.arg0)) {
// Device successfully created the socket
return;
}
addReverseTunnel(service: string, handler: AdbIncomingSocketHandler) {
this.#incomingSocketHandlers.set(service, handler);
const socket = this.#sockets.get(packet.arg1);
if (socket) {
// Device has received last `WRTE` to the socket
socket.ack();
return;
}
removeReverseTunnel(address: string) {
this.#incomingSocketHandlers.delete(address);
}
clearReverseTunnels() {
this.#incomingSocketHandlers.clear();
// 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 #handleOpen(packet: AdbPacketData) {
@ -240,12 +214,41 @@ export class AdbPacketDispatcher implements Closeable {
try {
await handler(controller.socket);
this.#sockets.set(localId, controller);
await this.sendPacket(AdbCommand.OK, localId, remoteId);
await this.sendPacket(AdbCommand.Okay, localId, remoteId);
} 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 (!handled) {
throw new Error(
`packet for \`${socket.service}\` not handled in 5 seconds`,
);
}
}),
(async () => {
await socket.enqueue(packet.payload);
await this.sendPacket(
AdbCommand.Okay,
packet.arg1,
packet.arg0,
);
handled = true;
})(),
]);
return;
}
async createSocket(service: string): Promise<AdbSocket> {
if (this.options.appendNullToServiceString) {
service += "\0";
@ -268,6 +271,18 @@ export class AdbPacketDispatcher implements Closeable {
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,
@ -306,7 +321,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();
@ -317,7 +332,7 @@ export class AdbPacketDispatcher implements Closeable {
#dispose() {
for (const socket of this.#sockets.values()) {
socket.dispose().catch(unreachable);
socket.dispose();
}
this.#disconnected.resolve();

View file

@ -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'
}

View file

@ -5,16 +5,15 @@ import type {
PushReadableStreamController,
ReadableStream,
WritableStream,
WritableStreamDefaultController,
} from "@yume-chan/stream-extra";
import {
ConsumableWritableStream,
DistributionStream,
DuplexStreamFactory,
PushReadableStream,
pipeFrom,
} from "@yume-chan/stream-extra";
import type { AdbSocket } from "../adb.js";
import { raceSignal } from "../server/index.js";
import type { AdbPacketDispatcher } from "./dispatcher.js";
import { AdbCommand } from "./packet.js";
@ -44,8 +43,6 @@ export class AdbDaemonSocketController
readonly localCreated!: boolean;
readonly service!: string;
#duplex: DuplexStreamFactory<Uint8Array, Consumable<Uint8Array>>;
#readable: ReadableStream<Uint8Array>;
#readableController!: PushReadableStreamController<Uint8Array>;
get readable() {
@ -53,16 +50,14 @@ export class AdbDaemonSocketController
}
#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;
@ -77,66 +72,44 @@ 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;
},
dispose: () => {
// Error out the pending writes
this.#writePromise?.reject(new Error("Socket closed"));
},
this.#readable = new PushReadableStream((controller) => {
this.#readableController = controller;
});
this.#readable = this.#duplex.wrapReadable(
new PushReadableStream(
(controller) => {
this.#readableController = controller;
this.writable = new ConsumableWritableStream<Uint8Array>({
start: (controller) => {
this.#writableController = controller;
},
{ highWaterMark: 0 },
),
);
this.writable = pipeFrom(
this.#duplex.createWritable(
new ConsumableWritableStream<Uint8Array>({
write: async (chunk) => {
// Wait for an ack packet
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
) {
this.#writePromise = new PromiseResolver();
await this.#dispatcher.sendPacket(
AdbCommand.Write,
this.localId,
this.remoteId,
chunk,
data.subarray(start, end),
);
await this.#writePromise.promise;
// Wait for ack packet
await raceSignal(
() => this.#writePromise!.promise,
controller.signal,
);
}
},
}),
),
new DistributionStream(this.#dispatcher.options.maxPayloadSize),
);
});
this.#socket = new AdbDaemonSocket(this);
}
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;
}
@ -149,11 +122,32 @@ export class AdbDaemonSocketController
}
async close(): Promise<void> {
await this.#duplex.close();
if (this.#closed) {
return;
}
this.#closed = true;
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();
}
}
@ -188,7 +182,7 @@ export class AdbDaemonSocket implements AdbDaemonSocketInfo, AdbSocket {
return this.#controller.writable;
}
get closed(): boolean {
get closed(): Promise<void> {
return this.#controller.closed;
}

View file

@ -154,7 +154,7 @@ export class BufferedReadableStream implements AsyncExactReadable {
}
}
cancel(reason?: unknown) {
return this.reader.cancel(reason);
async cancel(reason?: unknown) {
await this.reader.cancel(reason);
}
}

View file

@ -1,6 +1,10 @@
import { PromiseResolver } from "@yume-chan/async";
import type { QueuingStrategy, WritableStreamDefaultWriter } from "./stream.js";
import type {
QueuingStrategy,
WritableStreamDefaultController,
WritableStreamDefaultWriter,
} from "./stream.js";
import { ReadableStream, TransformStream, WritableStream } from "./stream.js";
interface Task {
@ -161,8 +165,13 @@ export class ConsumableReadableStream<T> extends ReadableStream<Consumable<T>> {
}
export interface ConsumableWritableStreamSink<T> {
start?(): void | PromiseLike<void>;
write?(chunk: T): void | PromiseLike<void>;
start?(
controller: WritableStreamDefaultController,
): void | PromiseLike<void>;
write?(
chunk: T,
controller: WritableStreamDefaultController,
): void | PromiseLike<void>;
abort?(reason: any): void | PromiseLike<void>;
close?(): void | PromiseLike<void>;
}
@ -196,11 +205,13 @@ export class ConsumableWritableStream<T> extends WritableStream<Consumable<T>> {
super(
{
start() {
return sink.start?.();
start(controller) {
return sink.start?.(controller);
},
async write(chunk) {
await chunk.tryConsume((value) => sink.write?.(value));
async write(chunk, controller) {
await chunk.tryConsume(
(value) => sink.write?.(value, controller),
);
chunk.consume();
},
abort(reason) {