fix(storage-plugin): preserve pending upload sockets across retries

packages/storage-plugin/lib/plugin.ts

@@ -4,6 +4,7 @@ import {fs, logger, tempDir, util} from '@appium/support';
 import type {AppiumServer} from '@appium/types';
 import {getResponseForW3CError} from 'appium/driver.js';
 import {BasePlugin} from 'appium/plugin.js';
+import AsyncLock from 'async-lock';
 import type {Express, Request, Response} from 'express';
 import {LRUCache} from 'lru-cache';
 import {WebSocketServer} from 'ws';
@@ -25,12 +26,7 @@ const WS_TTL_MS = 5 * 60 * 1000;
 const STORAGE_HANDLERS: Record<string, (req: Request, httpServer?: AppiumServer, basePath?: string) => Promise<any>> =
   {};
 const deprecatedRoutesLogged: Set<string> = new Set();
-const STORAGE_ADDITIONS_CACHE: LRUCache<string, () => any> = new LRUCache({
-  max: 20,
-  ttl: WS_TTL_MS,
-  ttlAutopurge: true,
-  dispose: (f: () => any) => f(),
-});
+const pendingUploads = new WeakMap<AppiumServer, {lock: AsyncLock; cache: LRUCache<string, () => void>}>();
 
 export class StoragePlugin extends BasePlugin {
   static async updateServer(expressApp: Express, httpServer: AppiumServer): Promise<void> {
@@ -93,7 +89,10 @@ STORAGE_HANDLERS.addStorageItem = async function addStorageItem(
     throw new Error('httpServer is required to add a storage item');
   }
   const itemOptions = requireValidItemOptions(parseRequestArgs(req, ['name', 'sha1']) as ItemOptions);
-  const [stream, events] = await prepareWebSockets(httpServer, itemOptions, basePath);
+  const {lock} = getPendingUploads(httpServer);
+  const [stream, events] = await lock.acquire(`${basePath}/${itemOptions.sha1}`, () =>
+    prepareWebSockets(httpServer, itemOptions, basePath),
+  );
   return {
     ws: {
       stream,
@@ -150,7 +149,7 @@ async function prepareWebSockets(
   const commonPathname = `${basePath}/add/${itemOptions.sha1}`;
   const streamPathname = `${commonPathname}/stream`;
   const eventsPathname = `${commonPathname}/events`;
-  if (!util.isEmpty(httpServer.getWebSocketHandlers(streamPathname))) {
+  if (Object.hasOwn(await httpServer.getWebSocketHandlers(streamPathname), streamPathname)) {
     return [streamPathname, eventsPathname];
   }
 
@@ -171,7 +170,8 @@ async function prepareWebSockets(
       signaler.removeAllListeners();
     }, 100);
   };
-  STORAGE_ADDITIONS_CACHE.set(itemOptions.sha1, streamDoneCallback);
+  const {cache} = getPendingUploads(httpServer);
+  cache.set(commonPathname, streamDoneCallback);
   eventsServer.on('connection', async (wsUpstream: WebSocket) => {
     signaler.on('status', (value) => wsUpstream.send(JSON.stringify(value)));
   });
@@ -191,7 +191,7 @@ async function prepareWebSockets(
       };
       log.debug(`Notifying about the successful addition of '${itemOptions.name}' to the server storage`);
       signaler.emit('status', successEvent);
-      STORAGE_ADDITIONS_CACHE.delete(itemOptions.sha1);
+      cache.delete(commonPathname);
     } catch (e) {
       log.debug(`Notifying about a failure while adding '${itemOptions.name}' to the server storage`);
       // in case of a failure we do not want to close the server yet
@@ -212,6 +212,23 @@ async function prepareWebSockets(
   return [streamPathname, eventsPathname];
 }
 
+function getPendingUploads(httpServer: AppiumServer) {
+  let pending = pendingUploads.get(httpServer);
+  if (!pending) {
+    pending = {
+      lock: new AsyncLock(),
+      cache: new LRUCache<string, () => void>({
+        max: 20,
+        ttl: WS_TTL_MS,
+        ttlAutopurge: true,
+        dispose: (cleanup) => cleanup(),
+      }),
+    };
+    pendingUploads.set(httpServer, pending);
+  }
+  return pending;
+}
+
 const getStorageSingleton = util.memoize(async () => {
   let storageRoot: string;
   let shouldPreserveRoot = false;

packages/storage-plugin/lib/storage.ts

@@ -15,7 +15,7 @@ import type {ItemOptions, StorageItem} from './types.js';
 const MAX_TASKS = 5;
 const TMP_EXT = '.filepart';
 const ADDITION_LOCK = new AsyncLock();
-const SHA1_HASH_LEN = 40;
+const SHA1_HASH_PATTERN = /^[a-f\d]{40}$/i;
 
 export class Storage {
   private readonly _root: string;
@@ -226,7 +226,7 @@ export class StorageArgumentError extends errors.InvalidArgumentError {}
  */
 export function requireValidItemOptions(opts: ItemOptions): ItemOptions {
   validateStorageItemName(opts.name);
-  if (opts.sha1?.length !== SHA1_HASH_LEN) {
+  if (typeof opts.sha1 !== 'string' || !SHA1_HASH_PATTERN.test(opts.sha1)) {
     throw new StorageArgumentError(
       `The provided hash value '${opts.sha1}' must be a valid SHA1 string, for ` +
         `example 'ccc963411b2621335657963322890305ebe96186'`,

packages/storage-plugin/test/unit/plugin.spec.ts

@@ -1,7 +1,7 @@
 import assert from 'node:assert/strict';
 import {describe, it} from 'node:test';
 
-import type {AppiumServer} from '@appium/types';
+import type {AppiumServer, WSServer} from '@appium/types';
 import type {Express, Request, Response} from 'express';
 
 import {StoragePlugin} from '../../lib/plugin.js';
@@ -15,14 +15,17 @@ interface CapturedResponse {
  * Registers the plugin routes against a stub app and returns a function which invokes the
  * handler bound to `routePath` with the given request body.
  */
-async function routeCaller(routePath: string, basePath = ''): Promise<(body: unknown) => Promise<CapturedResponse>> {
+async function routeCaller(
+  routePath: string,
+  basePath = '',
+  httpServer = {basePath} as AppiumServer,
+): Promise<(body: unknown) => Promise<CapturedResponse>> {
   const routes: Record<string, (req: Request, res: Response) => Promise<void>> = {};
   const register = (path: string, handler: (req: Request, res: Response) => Promise<void>) => {
     routes[path] = handler;
   };
   const app = {post: register, get: register} as unknown as Express;
-  // the request under test is rejected before the websocket setup runs, so this stub is never used
-  await StoragePlugin.updateServer(app, {basePath} as unknown as AppiumServer);
+  await StoragePlugin.updateServer(app, httpServer);
   const handler = routes[routePath];
   assert.ok(handler, `No handler was registered for ${routePath}`);
 
@@ -45,6 +48,60 @@ async function routeCaller(routePath: string, basePath = ''): Promise<(body: unk
 }
 
 describe('StoragePlugin routes', function () {
+  for (const sha1 of ['['.repeat(40), '/'.repeat(40), 'g'.repeat(40), Array(40).fill('a')]) {
+    it(`should reject invalid SHA1 input ${JSON.stringify(sha1)} before mounting sockets`, async function () {
+      const server = makeServer();
+      const callRoute = await routeCaller('/appium/storage/add', '', server);
+      const {status, body} = await callRoute({name: 'test.apk', sha1});
+      assert.equal(status, 400);
+      assert.equal(body.value.error, 'invalid argument');
+      assert.equal(server.registrations, 0);
+    });
+  }
+  for (const concurrent of [false, true]) {
+    it(`should reuse sockets for ${concurrent ? 'concurrent' : 'retried'} upload requests`, async function () {
+      const server = makeServer();
+      const callRoute = await routeCaller('/appium/storage/add', '', server);
+      const opts = {name: 'test.apk', sha1: 'a'.repeat(40)};
+      try {
+        let first: CapturedResponse;
+        let second: CapturedResponse;
+        if (concurrent) {
+          [first, second] = await Promise.all([callRoute(opts), callRoute(opts)]);
+        } else {
+          first = await callRoute(opts);
+          second = await callRoute(opts);
+        }
+        assert.equal(first.status, 200);
+        assert.deepEqual(first, second);
+        assert.equal(server.registrations, 2);
+        assert.equal(server.removals, 0);
+      } finally {
+        await closeSockets(server);
+      }
+    });
+  }
+
+  it('should keep uploads on different route prefixes and servers independent', async function () {
+    const firstServer = makeServer();
+    const secondServer = makeServer();
+    const opts = {name: 'test.apk', sha1: 'b'.repeat(40)};
+    try {
+      for (const server of [firstServer, secondServer]) {
+        for (const prefix of ['/appium/storage', '/storage']) {
+          const callRoute = await routeCaller(`${prefix}/add`, '', server);
+          assert.equal((await callRoute(opts)).status, 200);
+        }
+      }
+      assert.equal(firstServer.registrations, 4);
+      assert.equal(secondServer.registrations, 4);
+      assert.equal(firstServer.removals, 0);
+      assert.equal(secondServer.removals, 0);
+    } finally {
+      await closeSockets(firstServer);
+      await closeSockets(secondServer);
+    }
+  });
   it('should reject an invalid add name as an invalid argument', async function () {
     const callRoute = await routeCaller('/appium/storage/add');
     const {status, body} = await callRoute({
@@ -62,3 +119,33 @@ describe('StoragePlugin routes', function () {
     await assert.rejects(routeCaller('/appium/storage/add', '/wd/hub'));
   });
 });
+
+function makeServer(): AppiumServer & {registrations: number; removals: number} {
+  const mapping: Record<string, WSServer> = {};
+  const server = {
+    basePath: '',
+    webSocketsMapping: mapping,
+    registrations: 0,
+    removals: 0,
+    async getWebSocketHandlers(path: string) {
+      return Object.fromEntries(Object.entries(mapping).filter(([key]) => key.includes(path)));
+    },
+    async addWebSocketHandler(path: string, handler: WSServer) {
+      server.registrations++;
+      mapping[path] = handler;
+    },
+    async removeWebSocketHandler(path: string) {
+      server.removals++;
+      mapping[path]?.close();
+      delete mapping[path];
+      return true;
+    },
+  };
+  return server as unknown as AppiumServer & {registrations: number; removals: number};
+}
+
+async function closeSockets(server: AppiumServer): Promise<void> {
+  for (const path of Object.keys(server.webSocketsMapping)) {
+    await server.removeWebSocketHandler(path);
+  }
+}