import {createHash} from 'node:crypto'; import nativeFs from 'node:fs'; import path from 'node:path'; import type Stream from 'node:stream'; import {fs, timing, util} from '@appium/support'; import type {AppiumLogger} from '@appium/types'; import {errors} from 'appium/driver.js'; import AsyncLock from 'async-lock'; import {asyncmap} from 'asyncbox'; import type {Path} from 'path-scurry'; import type WebSocket from 'ws'; import type {ItemOptions, StorageItem} from './types.js'; const MAX_TASKS = 5; const TMP_EXT = '.filepart'; const ADDITION_LOCK = new AsyncLock(); const WS_SERVER_ERROR = 1011; const SHA1_HASH_LEN = 40; export class Storage { private readonly _root: string; private readonly _log: AppiumLogger; private readonly _shouldPreserveRoot: boolean; private readonly _shouldPreserveFiles: boolean; constructor(root: string, shouldPreserveRoot: boolean, shouldPreserveFiles: boolean, log: AppiumLogger) { this._root = root; this._log = log; this._shouldPreserveRoot = shouldPreserveRoot; this._shouldPreserveFiles = shouldPreserveFiles; } async list(): Promise { const items = (await this._listFiles()).filter((p) => !p.fullpath().endsWith(TMP_EXT)); if (util.isEmpty(items)) { return []; } const stats = await asyncmap(items, (item) => fs.stat(item.fullpath()), { concurrency: MAX_TASKS, }); return items.map((item, index) => ({ name: path.basename(item.fullpath()), path: item.fullpath(), size: stats[index].size, })); } async add(opts: ItemOptions, source: Stream | WebSocket): Promise { const {name} = requireValidItemOptions(opts); // toLowerCase is needed for case-insensitive server filesystems await ADDITION_LOCK.acquire(name.toLowerCase(), async () => { if (typeof (source as any).pipe === 'function') { await this._addFromStream(opts, source as Stream); } else { await this._addFromWebSocket(opts, source as WebSocket); } }); } async delete(name: string): Promise { if (name.toLowerCase().endsWith(TMP_EXT)) { return false; } const destinationPath = path.join(this._root, name); if (!(await fs.exists(destinationPath))) { return false; } await fs.rimraf(destinationPath); return true; } async reset(): Promise { if (!this._shouldPreserveRoot && !this._shouldPreserveFiles) { await fs.rimraf(this._root); } if (!(await fs.exists(this._root))) { await fs.mkdirp(this._root); return; } const files = (await this._listFiles()) .map((p) => p.fullpath()) .filter((fullPath) => !this._shouldPreserveFiles || path.basename(fullPath).toLowerCase().endsWith(TMP_EXT)); if (util.isEmpty(files)) { return; } await asyncmap(files, (fullPath) => fs.rimraf(fullPath), {concurrency: MAX_TASKS}); } cleanupSync(): void { this._log.debug(`Cleaning up the '${this._root}' server storage folder`); if (!this._shouldPreserveRoot && !this._shouldPreserveFiles) { fs.rimrafSync(this._root); return; } let itemNames: string[]; try { itemNames = nativeFs.readdirSync(this._root).filter((name) => !name.startsWith('.')); } catch (e) { const message = e instanceof Error ? e.message : String(e); this._log.warn( `Cannot list the '${this._root}' server storage folder. Original error: ${message}. ` + `Skipping the cleanup.`, ); return; } if (util.isEmpty(itemNames)) { if (!this._shouldPreserveRoot) { fs.rimrafSync(this._root); } return; } const matchedNames = itemNames.filter((name) => !this._shouldPreserveFiles || name.toLowerCase().endsWith(TMP_EXT)); for (const matchedName of matchedNames) { fs.rimrafSync(path.join(this._root, matchedName)); } const remainingNames = itemNames.filter((name) => !matchedNames.includes(name)); if (!this._shouldPreserveRoot && util.isEmpty(remainingNames)) { fs.rimrafSync(this._root); } } private async _listFiles(): Promise { const paths = (await fs.glob('*', { cwd: this._root, withFileTypes: true, })) as unknown as Path[]; return paths.filter((item) => item.isFile()); } private async _addFromStream(opts: ItemOptions, source: Stream): Promise { const {name} = opts; const fullPath = path.join(this._root, toTempName(name)); const timer = new timing.Timer().start(); const destination = fs.createWriteStream(fullPath); source.pipe(destination); try { await new Promise((resolve, reject) => { destination.once('finish', () => resolve(true)); source.once('error', reject); destination.once('error', reject); }); await this._finalizeItem(opts, timer, fullPath, await fs.hash(fullPath)); } catch (e) { await fs.rimraf(fullPath); throw e; } } private async _addFromWebSocket(opts: ItemOptions, source: WebSocket): Promise { const {name, sha1} = opts; const fullPath = path.join(this._root, toTempName(name)); const timer = new timing.Timer().start(); const destination = fs.createWriteStream(fullPath); const sha1sum = createHash('sha1'); let didDigestMatch = false; let recentDigest: string | null = null; try { await new Promise((resolve, reject) => { source.on('message', (data: WebSocket.RawData) => { if (didDigestMatch) { // ignore further chunks if hashes have already matched return; } destination.write(data, (e) => { if (e) { source.close(WS_SERVER_ERROR); reject(e); } }); sha1sum.update(data as any); recentDigest = sha1sum.copy().digest('hex'); if (recentDigest.toLowerCase() === sha1.toLowerCase()) { didDigestMatch = true; destination.close(() => resolve(true)); } }); source.once('close', () => { destination.close(() => resolve(true)); }); source.once('error', reject); destination.once('error', (e) => { source.close(WS_SERVER_ERROR); reject(e); }); }); await this._finalizeItem(opts, timer, fullPath, recentDigest ?? sha1sum.digest('hex')); } catch (e) { await fs.rimraf(fullPath); throw e; } } private async _finalizeItem( opts: ItemOptions, timer: timing.Timer, fullPath: string, actualHashDigest: string, ): Promise { const {name, sha1} = opts; this._log.info( `'${name}' has been added to the server storage within ` + `${timer.getDuration().asMilliSeconds}ms. Verifying hashes.`, ); if (actualHashDigest.toLowerCase() !== sha1.toLowerCase()) { throw new StorageArgumentError( `The actual SHA1 hash value '${actualHashDigest}' must be equal ` + `to the expected hash value of '${sha1}' for '${name}'`, ); } await fs.mv(fullPath, path.join(this._root, name)); } } export class StorageArgumentError extends errors.InvalidArgumentError {} /** * Validates storage item options and returns the same object when valid. * @param opts Candidate item options. */ export function requireValidItemOptions(opts: ItemOptions): ItemOptions { validateStorageItemName(opts.name); if (opts.sha1?.length !== SHA1_HASH_LEN) { throw new StorageArgumentError( `The provided hash value '${opts.sha1}' must be a valid SHA1 string, for ` + `example 'ccc963411b2621335657963322890305ebe96186'`, ); } return opts; } /** * Validate storage item name and throw if it is invalid. * @param name The name to validate. */ export function validateStorageItemName(name: string): void { if (util.isEmpty(name)) { throw new StorageArgumentError(`The provided file name '${name}' must not be empty`); } const sanitizedName = fs.sanitizeName(name, { replacement: '_', }); if (name !== sanitizedName) { throw new StorageArgumentError( `The provided name value '${name}' must be a valid file name. ` + `Did you mean '${sanitizedName}'?`, ); } } function toTempName(origName: string): string { return `${origName}${TMP_EXT}`; }