sha256 emscripten

Author: ngxson

packages/hub/.eslintignore

@@ -1 +1,2 @@
 dist
+sha256.js

packages/hub/src/index.ts

@@ -18,3 +18,5 @@ export type {
 	SpaceStage,
 } from "./types/public";
 export { HubApiError, InvalidApiResponseFormatError } from "./error";
+export { sha256 } from "./utils/sha256";
+export { WebBlob } from "./utils/WebBlob";

packages/hub/src/utils/sha256.spec.ts

@@ -0,0 +1,22 @@
+import { describe, it, expect } from "vitest";
+import { WebBlob } from "./WebBlob";
+import { sha256 } from "./sha256";
+
+describe("sha256", () => {
+	const resourceUrl = new URL(
+		"https://huggingface.co/ngxson/tinyllama_split_test/resolve/main/stories15M-q8_0-00001-of-00003.gguf"
+	);
+
+	it("Calculate hash in nodejs", async () => {
+		const blob = await WebBlob.create(resourceUrl, { cacheBelow: 0 });
+		const iterator = sha256(blob, { useWebWorker: true });
+		let res: IteratorResult<number, string>;
+		do {
+			res = await iterator.next();
+		} while (!res.done);
+		const sha = res.value;
+		expect(sha).toBe("8a3a74042ae05dda34985dff2d49adf2ee3f9d0fd73b03fa6cc4307c924e4040");
+	});
+
+	// TODO: how to test on browser (with / without web worker)
+});

packages/hub/src/utils/sha256.ts

@@ -2,31 +2,10 @@ import { eventToGenerator } from "./eventToGenerator";
 import { hexFromBytes } from "./hexFromBytes";
 import { isFrontend } from "./isFrontend";
 
-const webWorkerCode = `
-// Would prefer no CDN, but need a clever way to not burden the main file of the bundle
-importScripts("https://cdn.jsdelivr.net/npm/hash-wasm@4/dist/sha256.umd.min.js");
-
-const createSHA256 = hashwasm.createSHA256;
-
-self.addEventListener('message', async (event) => {
-	const { file } = event.data;
-	const sha256 = await createSHA256();
-	sha256.init();
-	const reader = file.stream().getReader();
-	const total = file.size;
-	let bytesDone = 0;
-	while (true) {
-		const { done, value } = await reader.read();
-		if (done) {
-			break;
-		}
-		sha256.update(value);
-		bytesDone += value.length;
-		postMessage({ progress: bytesDone / total });
-	}
-	postMessage({ sha256: sha256.digest('hex') });
-});
-`;
+async function getWebWorkerCode() {
+	const sha256Module = await import("../vendor/hash-wasm/sha256-wrapper");
+	return URL.createObjectURL(new Blob([sha256Module.createSHA256WorkerCode()]));
+}
 
 const pendingWorkers: Worker[] = [];
 const runningWorkers: Set<Worker> = new Set();
@@ -45,7 +24,7 @@ async function getWorker(poolSize?: number): Promise<Worker> {
 		}
 	}
 	if (!poolSize) {
-		const worker = new Worker(URL.createObjectURL(new Blob([webWorkerCode])));
+		const worker = new Worker(await getWebWorkerCode());
 		runningWorkers.add(worker);
 		return worker;
 	}
@@ -58,7 +37,7 @@ async function getWorker(poolSize?: number): Promise<Worker> {
 		await waitPromise;
 	}
 
-	const worker = new Worker(URL.createObjectURL(new Blob([webWorkerCode])));
+	const worker = new Worker(await getWebWorkerCode());
 	runningWorkers.add(worker);
 	return worker;
 }
@@ -147,7 +126,7 @@ export async function* sha256(
 			}
 		}
 		if (!wasmModule) {
-			wasmModule = await import("hash-wasm");
+			wasmModule = await import("../vendor/hash-wasm/sha256-wrapper");
 		}
 
 		const sha256 = await wasmModule.createSHA256();
@@ -184,4 +163,4 @@ export async function* sha256(
 // eslint-disable-next-line @typescript-eslint/consistent-type-imports
 let cryptoModule: typeof import("./sha256-node");
 // eslint-disable-next-line @typescript-eslint/consistent-type-imports
-let wasmModule: typeof import("hash-wasm");
+let wasmModule: typeof import("../vendor/hash-wasm/sha256-wrapper");

packages/hub/src/vendor/hash-wasm/build.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+CURRENT_PATH="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
+cd $CURRENT_PATH
+
+# Clean up
+docker kill hash-wasm-builder
+docker rm hash-wasm-builder
+
+# Start container
+docker run -it -d --name hash-wasm-builder emscripten/emsdk:3.1.55 bash
+
+# Copy & compile
+docker exec hash-wasm-builder bash -c "mkdir /source"
+docker cp ./sha256.c hash-wasm-builder:/source
+docker exec hash-wasm-builder bash -c "\
+  cd /source && \
+  emcc sha256.c -o sha256.js -msimd128 -sSINGLE_FILE -sMODULARIZE=1 -sENVIRONMENT=web,worker -sEXPORTED_FUNCTIONS=_Hash_Init,_Hash_Update,_Hash_Final,_GetBufferPtr -sFILESYSTEM=0 -fno-rtti -fno-exceptions -O1 -sMODULARIZE=1 \
+  "
+docker cp hash-wasm-builder:/source/sha256.js .
+
+# Clean up
+docker kill hash-wasm-builder
+docker rm hash-wasm-builder

packages/hub/src/vendor/hash-wasm/sha256-wrapper.ts

@@ -0,0 +1,66 @@
+import WasmModule from "./sha256";
+
+const BUFFER_MAX_SIZE = 8 * 1024 * 1024;
+
+export async function createSHA256(isInsideWorker = false): Promise<{
+	init(): void,
+	update(data: Uint8Array): void,
+	digest(method: "hex"): string,
+}> {
+	const wasm: Awaited<ReturnType<typeof WasmModule>> = isInsideWorker
+		? // @ts-expect-error WasmModule will be populated inside self object
+		  await self["WasmModule"]()
+		: await WasmModule();
+		const heap = wasm.HEAPU8.subarray(wasm._GetBufferPtr());
+	return {
+		init() {
+			wasm._Hash_Init(256);
+		},
+		update(data: Uint8Array) {
+			let byteUsed = 0;
+			while (byteUsed < data.byteLength) {
+				const bytesLeft = data.byteLength - byteUsed;
+				const length = bytesLeft < BUFFER_MAX_SIZE ? bytesLeft : BUFFER_MAX_SIZE;
+				heap.set(data.subarray(byteUsed, length));
+				wasm._Hash_Update(length);
+				byteUsed += length;
+			}
+		},
+		digest(method: "hex") {
+			if (method !== "hex") {
+				throw new Error("Only digest hex is supported");
+			}
+			wasm._Hash_Final();
+			const result = Array.from(heap.slice(0, 32));
+			return result.map((b) => b.toString(16).padStart(2, "0")).join("");
+		},
+	};
+}
+
+export function createSHA256WorkerCode(): string {
+	return `
+    const _scriptDir = "";
+    self.WasmModule = ${WasmModule.toString()};
+
+    ${createSHA256.toString()};
+
+    self.addEventListener('message', async (event) => {
+      const { file } = event.data;
+      const sha256 = await createSHA256(true);
+      sha256.init();
+      const reader = file.stream().getReader();
+      const total = file.size;
+      let bytesDone = 0;
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          break;
+        }
+        sha256.update(value);
+        bytesDone += value.length;
+        postMessage({ progress: bytesDone / total });
+      }
+      postMessage({ sha256: sha256.digest('hex') });
+    });
+  `;
+}