Merge pull request #20 from bakkot/prevent-excessive-requests

Allow bounding total requested data

Author: phiresky

README.md

@@ -64,9 +64,13 @@ const config = {
   configUrl: "/foo/bar/config.json"
 }
 
+
+let maxBytesToRead = 10 * 1024 * 1024;
 const worker = await createDbWorker(
   [config],
-  workerUrl.toString(), wasmUrl.toString()
+  workerUrl.toString(),
+  wasmUrl.toString(),
+  maxBytesToRead // optional, defaults to Infinity
 );
 // you can also pass multiple config objects which can then be used as separate database schemas with `ATTACH virtualFilename as schemaname`, where virtualFilename is also set in the config object.
 
@@ -75,6 +79,12 @@ const worker = await createDbWorker(
 
 const result = await worker.db.exec(`select * from table where id = ?`, [123]);
 
+// worker.worker.bytesRead is a Promise for the number of bytes read by the worker.
+// if a request would cause it to exceed maxBytesToRead, that request will throw a SQLite disk I/O error.
+console.log(await worker.worker.bytesRead);
+
+// you can reset bytesRead by assigning to it:
+worker.worker.bytesRead = 0;
 ```
 
 ## Cachebusting
@@ -186,4 +196,4 @@ cd sql.js
 yarn build
 cd ..
 yarn build
-```
+```

sql.js/package-lock.json

@@ -31,14 +31,17 @@ export interface WorkerHttpvfs {
 export async function createDbWorker(
   configs: SplitFileConfig[],
   workerUrl: string,
-  wasmUrl: string
+  wasmUrl: string,
+  maxBytesToRead: number = Infinity
 ): Promise<WorkerHttpvfs> {
   const worker: Worker = new Worker(workerUrl);
   const sqlite = Comlink.wrap<SqliteComlinkMod>(worker);
 
   const db = ((await sqlite.SplitFileHttpDatabase(
     wasmUrl,
-    configs
+    configs,
+    undefined,
+    maxBytesToRead
   )) as unknown) as Comlink.Remote<LazyHttpDatabase>;
 
   worker.addEventListener("message", handleAsyncRequestFromWorkerThread);

src/db.ts

@@ -8,6 +8,8 @@ export type RangeMapper = (
   toByte: number
 ) => { url: string; fromByte: number; toByte: number };
 
+export type RequestLimiter = (bytes: number) => void;
+
 export type LazyFileConfig = {
   /** function to map a read request to an url with read request  */
   rangeMapper: RangeMapper;
@@ -21,6 +23,8 @@ export type LazyFileConfig = {
   maxReadSpeed?: number;
   /** if true, log all read pages into the `readPages` field for debugging */
   logPageReads?: boolean;
+  /** if defined, this is called once per request and passed the number of bytes about to be requested **/
+  requestLimiter?: RequestLimiter;
 };
 export type PageReadLog = {
   pageno: number;
@@ -46,6 +50,7 @@ export class LazyUint8Array {
   private readonly maxSpeed: number;
   private readonly maxReadHeads: number;
   private readonly logPageReads: boolean;
+  private readonly requestLimiter: RequestLimiter;
 
   constructor(config: LazyFileConfig) {
     this._chunkSize = config.requestChunkSize;
@@ -58,6 +63,7 @@ export class LazyUint8Array {
     if (config.fileLength) {
       this._length = config.fileLength;
     }
+    this.requestLimiter = config.requestLimiter == null ? ((ignored) => {}) : config.requestLimiter;
   }
   /**
    * efficiently copy the range [start, start + length) from the http file into the
@@ -228,6 +234,7 @@ export class LazyUint8Array {
         absoluteFrom / 1024
       } KiB]`
     );
+    this.requestLimiter(absoluteTo - absoluteFrom);
     this.totalFetchedBytes += absoluteTo - absoluteFrom;
     this.totalRequests++;
     if (absoluteFrom > absoluteTo)

src/lazyFile.ts

@@ -127,17 +127,33 @@ const mod = {
   db: null as null | LazyHttpDatabase,
   inited: false,
   sqljs: null as null | Promise<any>,
+  bytesRead: 0,
   async SplitFileHttpDatabase(
     wasmUrl: string,
     configs: SplitFileConfig[],
-    mainVirtualFilename?: string
+    mainVirtualFilename?: string,
+    maxBytesToRead: number = Infinity,
   ): Promise<LazyHttpDatabase> {
     if (this.inited) throw Error(`sorry, only one db is supported right now`);
     this.inited = true;
     if (!this.sqljs) {
       this.sqljs = init(wasmUrl);
     }
     const sql = await this.sqljs;
+
+    this.bytesRead = 0;
+    let requestLimiter = (bytes: number) => {
+      if (this.bytesRead + bytes > maxBytesToRead) {
+        this.bytesRead = 0;
+        // I couldn't figure out how to get ERRNO_CODES included
+        // so just hardcode the actual value
+        // https://github.com/emscripten-core/emscripten/blob/565fb3651ed185078df1a13b8edbcb6b2192f29e/system/include/wasi/api.h#L146
+        // https://github.com/emscripten-core/emscripten/blob/565fb3651ed185078df1a13b8edbcb6b2192f29e/system/lib/libc/musl/arch/emscripten/bits/errno.h#L13
+        throw new sql.FS.ErrnoError(6 /* EAGAIN */);
+      }
+      this.bytesRead += bytes;
+    };
+
     const lazyFiles = new Map();
     const hydratedConfigs = await fetchConfigs(configs);
     let mainFileConfig;
@@ -182,7 +198,8 @@ const mod = {
             ? config.databaseLengthBytes
             : undefined,
         logPageReads: true,
-        maxReadHeads: 3
+        maxReadHeads: 3,
+        requestLimiter
       });
       lazyFiles.set(filename, lazyFile);
     }
@@ -197,7 +214,7 @@ const mod = {
           `Chunk size does not match page size: pragma page_size = ${pageSize} but chunkSize = ${mainFileConfig.requestChunkSize}`
         );
     }
-  
+
     this.db.lazyFiles = lazyFiles;
     this.db.create_vtab(SeriesVtab);
     this.db.query = (...args) => toObjects(this.db!.exec(...args));