Re-implement R2 gateway using new storage system

Closes DEVX-592

Author: mrbbot

packages/tre/src/plugins/r2/gateway.ts

@@ -1,108 +1,64 @@
-import assert from "assert";
-import crypto from "crypto";
-import { TextEncoder } from "util";
+import { Blob } from "buffer";
+import { ReadableStream, TransformStream } from "stream/web";
 import type { R2StringChecksums } from "@cloudflare/workers-types/experimental";
 import { R2Objects } from "./gateway";
 import {
-  BASE64_REGEXP,
   HEX_REGEXP,
+  ObjectRow,
   R2HeadResponse,
   R2HttpFields,
   R2Range,
 } from "./schemas";
 
-const encoder = new TextEncoder();
-
-export interface R2ObjectMetadata {
-  // The object’s key.
-  key: string;
-  // Random unique string associated with a specific upload of a key.
-  version: string;
-  // Size of the object in bytes.
-  size: number;
-  // The etag associated with the object upload.
-  etag: string;
-  // The object's etag, in quotes to be returned as a header.
-  httpEtag: string;
-  // The time the object was uploaded.
-  uploaded: number;
-  // Various HTTP headers associated with the object. Refer to HTTP Metadata:
-  // https://developers.cloudflare.com/r2/runtime-apis/#http-metadata.
-  httpMetadata: R2HttpFields;
-  // A map of custom, user-defined metadata associated with the object.
-  customMetadata: Record<string, string>;
-  // If a GET request was made with a range option, this will be added
-  range?: R2Range;
-  // Hashes used to check the received object’s integrity. At most one can be
-  // specified.
-  checksums?: R2StringChecksums;
-}
-
 export interface EncodedMetadata {
   metadataSize: number;
-  value: Uint8Array;
+  value: ReadableStream<Uint8Array>;
 }
 
-export function createVersion(): string {
-  return crypto.randomBytes(16).toString("hex");
-}
-
-/**
- * R2Object is created when you PUT an object into an R2 bucket.
- * R2Object represents the metadata of an object based on the information
- * provided by the uploader. Every object that you PUT into an R2 bucket
- * will have an R2Object created.
- */
-export class R2Object implements R2ObjectMetadata {
+export class R2Object {
   readonly key: string;
   readonly version: string;
   readonly size: number;
   readonly etag: string;
-  readonly httpEtag: string;
   readonly uploaded: number;
   readonly httpMetadata: R2HttpFields;
   readonly customMetadata: Record<string, string>;
   readonly range?: R2Range;
   readonly checksums: R2StringChecksums;
 
-  constructor(metadata: R2ObjectMetadata) {
-    this.key = metadata.key;
-    this.version = metadata.version;
-    this.size = metadata.size;
-    this.etag = metadata.etag;
-    this.httpEtag = metadata.httpEtag;
-    this.uploaded = metadata.uploaded;
-    this.httpMetadata = metadata.httpMetadata;
-    this.customMetadata = metadata.customMetadata;
-    this.range = metadata.range;
+  constructor(row: Omit<ObjectRow, "blob_id">, range?: R2Range) {
+    this.key = row.key;
+    this.version = row.version;
+    this.size = row.size;
+    this.etag = row.etag;
+    this.uploaded = row.uploaded;
+    this.httpMetadata = JSON.parse(row.http_metadata);
+    this.customMetadata = JSON.parse(row.custom_metadata);
+    this.range = range;
 
     // For non-multipart uploads, we always need to store an MD5 hash in
-    // `checksums`, but never explicitly stored one. Luckily, `R2Bucket#put()`
-    // always makes `etag` an MD5 hash.
-    const checksums: R2StringChecksums = { ...metadata.checksums };
-    const etag = metadata.etag;
-    if (etag.length === 32 && HEX_REGEXP.test(etag)) {
-      checksums.md5 = metadata.etag;
-    } else if (etag.length === 24 && BASE64_REGEXP.test(etag)) {
-      // TODO: remove this when we switch underlying storage mechanisms
-      // Previous versions of Miniflare 3 base64 encoded `etag` instead
-      checksums.md5 = Buffer.from(etag, "base64").toString("hex");
-    } else {
-      assert.fail("Expected `etag` to be an MD5 hash");
+    // `checksums`. To avoid data duplication, we just use `etag` for this.
+    const checksums: R2StringChecksums = JSON.parse(row.checksums);
+    if (this.etag.length === 32 && HEX_REGEXP.test(this.etag)) {
+      checksums.md5 = row.etag;
     }
     this.checksums = checksums;
   }
 
   // Format for return to the Workers Runtime
   #rawProperties(): R2HeadResponse {
     return {
-      ...this,
       name: this.key,
+      version: this.version,
+      size: this.size,
+      etag: this.etag,
+      uploaded: this.uploaded,
       httpFields: this.httpMetadata,
       customFields: Object.entries(this.customMetadata).map(([k, v]) => ({
         k,
         v,
       })),
+      range: this.range,
       checksums: {
         0: this.checksums.md5,
         1: this.checksums.sha1,
@@ -115,36 +71,38 @@ export class R2Object implements R2ObjectMetadata {
 
   encode(): EncodedMetadata {
     const json = JSON.stringify(this.#rawProperties());
-    const bytes = encoder.encode(json);
-    return { metadataSize: bytes.length, value: bytes };
+    const blob = new Blob([json]);
+    return { metadataSize: blob.size, value: blob.stream() };
   }
 
   static encodeMultiple(objects: R2Objects): EncodedMetadata {
     const json = JSON.stringify({
       ...objects,
       objects: objects.objects.map((o) => o.#rawProperties()),
     });
-    const bytes = encoder.encode(json);
-    return { metadataSize: bytes.length, value: bytes };
+    const blob = new Blob([json]);
+    return { metadataSize: blob.size, value: blob.stream() };
   }
 }
 
 export class R2ObjectBody extends R2Object {
-  readonly body: Uint8Array;
-
-  constructor(metadata: R2ObjectMetadata, body: Uint8Array) {
-    super(metadata);
-    this.body = body;
+  constructor(
+    metadata: Omit<ObjectRow, "blob_id">,
+    readonly body: ReadableStream<Uint8Array>,
+    range?: R2Range
+  ) {
+    super(metadata, range);
   }
 
   encode(): EncodedMetadata {
     const { metadataSize, value: metadata } = super.encode();
-    const merged = new Uint8Array(metadataSize + this.body.length);
-    merged.set(metadata);
-    merged.set(this.body, metadataSize);
+    const identity = new TransformStream();
+    void metadata
+      .pipeTo(identity.writable, { preventClose: true })
+      .then(() => this.body.pipeTo(identity.writable));
     return {
       metadataSize: metadataSize,
-      value: merged,
+      value: identity.readable,
     };
   }
 }

packages/tre/src/plugins/r2/r2Object.ts

@@ -1,4 +1,5 @@
-import { TextDecoder } from "util";
+import assert from "assert";
+import { ReadableStream, TransformStream } from "stream/web";
 import { Request, Response } from "../../http";
 import {
   CfHeader,
@@ -13,23 +14,46 @@ import { R2Gateway, R2Objects } from "./gateway";
 import { EncodedMetadata, R2Object, R2ObjectBody } from "./r2Object";
 import { R2BindingRequestSchema } from "./schemas";
 
-const decoder = new TextDecoder();
-
 async function decodeMetadata(req: Request) {
-  const bytes = await req.arrayBuffer();
-
   const metadataSize = Number(req.headers.get(CfHeader.MetadataSize));
-  if (Number.isNaN(metadataSize)) {
-    throw new InvalidMetadata();
+  if (Number.isNaN(metadataSize)) throw new InvalidMetadata();
+
+  assert(req.body !== null);
+  const body = req.body as ReadableStream<Uint8Array>;
+
+  // Read just metadata from body stream
+  const chunks: Uint8Array[] = [];
+  let chunksLength = 0;
+  for await (const chunk of body.values({ preventCancel: true })) {
+    chunks.push(chunk);
+    chunksLength += chunk.byteLength;
+    // Once we've read enough bytes, stop
+    if (chunksLength >= metadataSize) break;
   }
+  // If we read the entire stream without enough bytes for metadata, throw
+  if (chunksLength < metadataSize) throw new InvalidMetadata();
+  const atLeastMetadata = Buffer.concat(chunks, chunksLength);
+  const metadataJson = atLeastMetadata.subarray(0, metadataSize).toString();
+  const metadata = R2BindingRequestSchema.parse(JSON.parse(metadataJson));
 
-  const [metadataBytes, value] = [
-    bytes.slice(0, metadataSize),
-    bytes.slice(metadataSize),
-  ];
-  const metadataText = decoder.decode(metadataBytes);
-  const metadata = R2BindingRequestSchema.parse(JSON.parse(metadataText));
-  return { metadata, value: new Uint8Array(value) };
+  let value = body;
+  // If we read some value when reading metadata (quite likely), create a new
+  // stream, write the bit we read, then write the rest of the body stream
+  if (chunksLength > metadataSize) {
+    const identity = new TransformStream();
+    const writer = identity.writable.getWriter();
+    // The promise returned by `writer.write()` will only resolve once the chunk
+    // is read, which won't be until after this function returns, so we can't
+    // use `await` here
+    void writer.write(atLeastMetadata.subarray(metadataSize)).then(() => {
+      // Release the writer without closing the stream
+      writer.releaseLock();
+      return body.pipeTo(identity.writable);
+    });
+    value = identity.readable;
+  }
+
+  return { metadata, metadataSize, value };
 }
 function decodeHeaderMetadata(req: Request) {
   const header = req.headers.get(CfHeader.Request);
@@ -80,7 +104,7 @@ export class R2Router extends Router<R2Gateway> {
 
   @PUT("/:bucket")
   put: RouteHandler<R2Params> = async (req, params) => {
-    const { metadata, value } = await decodeMetadata(req);
+    const { metadata, metadataSize, value } = await decodeMetadata(req);
     const persist = decodePersist(req.headers);
     const gateway = this.gatewayFactory.get(params.bucket, persist);
 
@@ -90,7 +114,18 @@ export class R2Router extends Router<R2Gateway> {
       );
       return new Response();
     } else if (metadata.method === "put") {
-      const result = await gateway.put(metadata.object, value, metadata);
+      const contentLength = Number(req.headers.get("Content-Length"));
+      // `workerd` requires a known value size for R2 put requests:
+      // - https://github.com/cloudflare/workerd/blob/e3479895a2ace28e4fd5f1399cea4c92291966ab/src/workerd/api/r2-rpc.c%2B%2B#L154-L156
+      // - https://github.com/cloudflare/workerd/blob/e3479895a2ace28e4fd5f1399cea4c92291966ab/src/workerd/api/r2-rpc.c%2B%2B#L188-L189
+      assert(!isNaN(contentLength));
+      const valueSize = contentLength - metadataSize;
+      const result = await gateway.put(
+        metadata.object,
+        value,
+        valueSize,
+        metadata
+      );
       return encodeResult(result);
     } else {
       throw new InternalError(); // Unknown method: should never be reached

packages/tre/src/plugins/r2/router.ts

@@ -1,5 +1,30 @@
 import { z } from "zod";
 
+export interface ObjectRow {
+  key: string;
+  blob_id: string | null; // null if multipart
+  version: string;
+  size: number; // total size of object (all parts) in bytes
+  etag: string; // hex MD5 hash if not multipart
+  uploaded: number; // milliseconds since unix epoch
+  checksums: string; // JSON-serialised `R2StringChecksums` (workers-types)
+  http_metadata: string; // JSON-serialised `R2HTTPMetadata` (workers-types)
+  custom_metadata: string; // JSON-serialised user-defined metadata
+}
+export const SQL_SCHEMA = `
+CREATE TABLE IF NOT EXISTS _mf_objects (
+    key TEXT PRIMARY KEY,
+    blob_id TEXT,
+    version TEXT NOT NULL,
+    size INTEGER NOT NULL,
+    etag TEXT NOT NULL,
+    uploaded INTEGER NOT NULL,
+    checksums TEXT NOT NULL,
+    http_metadata TEXT NOT NULL,
+    custom_metadata TEXT NOT NULL
+);
+`;
+
 // https://github.com/cloudflare/workerd/blob/4290f9717bc94647d9c8afd29602cdac97fdff1b/src/workerd/api/r2-api.capnp
 
 export const HEX_REGEXP = /^[0-9a-f]*$/i;
@@ -193,6 +218,11 @@ export const R2BindingRequestSchema = z.union([
   R2DeleteRequestSchema,
 ]);
 
+export type OmitRequest<T> = Omit<T, "method" | "object">;
+export type R2GetOptions = OmitRequest<z.infer<typeof R2GetRequestSchema>>;
+export type R2PutOptions = OmitRequest<z.infer<typeof R2PutRequestSchema>>;
+export type R2ListOptions = OmitRequest<z.infer<typeof R2ListRequestSchema>>;
+
 export interface R2ErrorResponse {
   version: number;
   v4code: number;

packages/tre/src/plugins/r2/schemas.ts

@@ -1,15 +1,18 @@
-import crypto from "crypto";
+import assert from "assert";
 import { R2StringChecksums } from "@cloudflare/workers-types/experimental";
+import { InclusiveRange } from "../../storage2";
+import { _parseRanges } from "../shared";
 import {
   BadDigest,
   EntityTooLarge,
   InvalidMaxKeys,
   InvalidObjectName,
+  InvalidRange,
   MetadataTooLarge,
   PreconditionFailed,
 } from "./errors";
-import { R2Object, R2ObjectMetadata } from "./r2Object";
-import { R2Conditional } from "./schemas";
+import { R2Object } from "./r2Object";
+import { R2Conditional, R2GetOptions } from "./schemas";
 
 export const MAX_LIST_KEYS = 1_000;
 const MAX_KEY_SIZE = 1024;
@@ -28,7 +31,7 @@ function truncateToSeconds(ms: number) {
 /** @internal */
 export function _testR2Conditional(
   cond: R2Conditional,
-  metadata?: Pick<R2ObjectMetadata, "etag" | "uploaded">
+  metadata?: Pick<R2Object, "etag" | "uploaded">
 ): boolean {
   // Adapted from internal R2 gateway implementation.
   // See also https://datatracker.ietf.org/doc/html/rfc7232#section-6.
@@ -79,12 +82,14 @@ function serialisedLength(x: string) {
 }
 
 export class Validator {
-  hash(value: Uint8Array, hashes: R2Hashes): R2StringChecksums {
+  hash(digests: Map<string, Buffer>, hashes: R2Hashes): R2StringChecksums {
     const checksums: R2StringChecksums = {};
     for (const { name, field } of R2_HASH_ALGORITHMS) {
       const providedHash = hashes[field];
       if (providedHash !== undefined) {
-        const computedHash = crypto.createHash(field).update(value).digest();
+        const computedHash = digests.get(field);
+        // Should've computed all required digests
+        assert(computedHash !== undefined);
         if (!providedHash.equals(computedHash)) {
           throw new BadDigest(name, providedHash, computedHash);
         }
@@ -96,17 +101,53 @@ export class Validator {
     return checksums;
   }
 
-  condition(meta?: R2Object, onlyIf?: R2Conditional): Validator {
+  condition(
+    meta?: R2Object | Pick<R2Object, "etag" | "uploaded">,
+    onlyIf?: R2Conditional
+  ): Validator {
     if (onlyIf !== undefined && !_testR2Conditional(onlyIf, meta)) {
       let error = new PreconditionFailed();
-      if (meta !== undefined) error = error.attach(meta);
+      if (meta instanceof R2Object) error = error.attach(meta);
       throw error;
     }
     return this;
   }
 
-  size(value: Uint8Array): Validator {
-    if (value.byteLength > MAX_VALUE_SIZE) {
+  range(
+    options: Pick<R2GetOptions, "rangeHeader" | "range">,
+    size: number
+  ): InclusiveRange | undefined {
+    if (options.rangeHeader !== undefined) {
+      const ranges = _parseRanges(options.rangeHeader, size);
+      // If the header contained a single range, use it. Otherwise, if the
+      // header was invalid, or contained multiple ranges, just return the full
+      // response.
+      if (ranges?.length === 1) {
+        const [start, end] = ranges[0];
+        return { start, end };
+      }
+    } else if (options.range !== undefined) {
+      let { offset, length, suffix } = options.range;
+      // Eliminate suffix is specified
+      if (suffix !== undefined) {
+        if (suffix <= 0) throw new InvalidRange();
+        if (suffix > size) suffix = size;
+        offset = size - suffix;
+        length = suffix;
+      }
+      // Validate offset and length
+      if (offset === undefined) offset = 0;
+      if (length === undefined) length = size - offset;
+      if (offset < 0 || offset > size || length <= 0) throw new InvalidRange();
+      // Clamp length to maximum
+      if (offset + length > size) length = size - offset;
+      // Convert to inclusive range
+      return { start: offset, end: offset + length - 1 }; // TODO: check -1
+    }
+  }
+
+  size(size: number): Validator {
+    if (size > MAX_VALUE_SIZE) {
       throw new EntityTooLarge();
     }
     return this;

packages/tre/src/plugins/r2/validator.ts

@@ -19,3 +19,8 @@ export type TypedDatabase = Omit<Database, "prepare"> & {
     ? TypedStatement<Params, SingleResult>
     : TypedStatement<[Params], SingleResult>;
 };
+
+export function escapeLike(prefix: string) {
+  // Prefix all instances of `\`, `_` and `%` with `\`
+  return prefix.replace(/[\\_%]/g, "\\$&");
+}

packages/tre/src/storage2/sql.ts

@@ -36,12 +36,7 @@ import {
 } from "@miniflare/tre";
 import { Macro, ThrowsExpectation } from "ava";
 import { z } from "zod";
-import {
-  MiniflareTestContext,
-  miniflareTest,
-  useTmp,
-  utf8Decode,
-} from "../../test-shared";
+import { MiniflareTestContext, miniflareTest, useTmp } from "../../test-shared";
 import { isWithin } from "../../test-shared/asserts";
 
 const WITHIN_EPSILON = 10_000;
@@ -967,9 +962,7 @@ test(
   {
     keys: ["key1", "key2", "key3", "key4"],
     options: { startAfter: "key1", limit: 2 },
-    // TODO(soon): this should be `[["key2", "key3"], ["key4"]]`, see comment in
-    //  `gateway.ts` for details, we'll fix this with the new storage system
-    pages: [["key2"], ["key3", "key4"]],
+    pages: [["key2", "key3"], ["key4"]],
   }
 );
 test(
@@ -1204,16 +1197,20 @@ test.serial("operations persist stored data", async (t) => {
 
   // Create new temporary file-system persistence directory
   const tmp = await useTmp(t);
-  const storage = new FileStorage(path.join(tmp, "bucket"));
+  const legacyStorage = new FileStorage(path.join(tmp, "bucket"));
+  const newStorage = legacyStorage.getNewStorage();
 
   // Set option, then reset after test
   await t.context.setOptions({ ...opts, r2Persist: tmp });
   t.teardown(() => t.context.setOptions(opts));
 
   // Check put respects persist
   await r2.put("key", "value");
-  const stored = await storage.get(`${ns}key`);
-  t.is(utf8Decode(stored?.value), "value");
+  const stmtListByNs = newStorage.db.prepare<{ ns: string }, { key: string }>(
+    "SELECT key FROM _mf_objects WHERE key LIKE :ns || '%'"
+  );
+  let stored = stmtListByNs.all({ ns });
+  t.deepEqual(stored, [{ key: `${ns}key` }]);
 
   // Check head respects persist
   const object = await r2.head("key");
@@ -1230,5 +1227,8 @@ test.serial("operations persist stored data", async (t) => {
 
   // Check delete respects persist
   await r2.delete("key");
-  t.false(await storage.has(`${ns}key`));
+  stored = stmtListByNs.all({ ns });
+  t.deepEqual(stored, []);
 });
+
+// TODO: add tests for empty key

packages/tre/test/plugins/r2/index.spec.ts

@@ -1,5 +1,5 @@
 import { R2Conditional } from "@cloudflare/workers-types/experimental";
-import { R2ObjectMetadata, _testR2Conditional } from "@miniflare/tre";
+import { R2Object, _testR2Conditional } from "@miniflare/tre";
 import test from "ava";
 
 test("testR2Conditional: matches various conditions", (t) => {
@@ -11,7 +11,7 @@ test("testR2Conditional: matches various conditions", (t) => {
   const pastDate = new Date(uploadedDate.getTime() - 30_000);
   const futureDate = new Date(uploadedDate.getTime() + 30_000);
 
-  const metadata: Pick<R2ObjectMetadata, "etag" | "uploaded"> = {
+  const metadata: Pick<R2Object, "etag" | "uploaded"> = {
     etag,
     uploaded: uploadedDate.getTime(),
   };