| 1 | import type { Env } from "../env.js"; |
| 2 | import { Sha256 } from "../lib/sha256.js"; |
| 3 | import { tailKey, type UploadRow } from "./store.js"; |
| 4 | |
| 5 | /** |
| 6 | * Blob uploads stream into R2. Small blobs become one object; bigger ones become a |
| 7 | * multipart upload with fixed-size parts (R2 wants every part but the last to be the |
| 8 | * same size). Bytes that don't fill a part yet wait in a tail object until the next |
| 9 | * chunk or the final PUT. The object key is the upload's, not the digest's: R2 can't |
| 10 | * rename, so the blob index in D1 records where each digest's bytes are. |
| 11 | */ |
| 12 | export const PART_SIZE = 10 * 1024 * 1024; |
| 13 | |
| 14 | export class UploadError extends Error {} |
| 15 | |
| 16 | const hex = (b: ArrayBuffer) => [...new Uint8Array(b)].map((x) => x.toString(16).padStart(2, "0")).join(""); |
| 17 | |
| 18 | /** |
| 19 | * Appends a request body to an upload, hashing it on the way through. The hash's progress |
| 20 | * is kept with the upload, so a chunked upload's last request doesn't read the whole blob |
| 21 | * back to hash it (for a big layer that took longer than clients wait for a response). |
| 22 | * Uploads started before the state was kept are hashed when they finish. |
| 23 | */ |
| 24 | export async function ingest(env: Env, u: UploadRow, body: ReadableStream<Uint8Array> | null, final: boolean, limit?: { bytes: number; error: () => Error }): Promise<UploadRow> { |
| 25 | const bucket = env.BUCKET; |
| 26 | const parts = JSON.parse(u.parts) as R2UploadedPart[]; |
| 27 | let mp = u.multipart_id ? bucket.resumeMultipartUpload(u.r2_key, u.multipart_id) : null; |
| 28 | let buf = new Uint8Array(PART_SIZE); |
| 29 | let fill = 0; |
| 30 | |
| 31 | // Bring back the bytes that are not in a finished part yet. |
| 32 | const pending = mp ? u.tail : u.size; |
| 33 | if (pending > 0) { |
| 34 | const obj = await bucket.get(mp ? tailKey(u.id) : u.r2_key); |
| 35 | if (!obj) throw new UploadError("the upload's data is missing; start the upload again"); |
| 36 | const bytes = new Uint8Array(await obj.arrayBuffer()); |
| 37 | buf.set(bytes); |
| 38 | fill = bytes.length; |
| 39 | } |
| 40 | |
| 41 | const hasher = u.size === 0 ? new Sha256() : u.hash_state ? Sha256.restore(u.hash_state) : null; |
| 42 | let added = 0; |
| 43 | const flush = async () => { |
| 44 | mp ??= await bucket.createMultipartUpload(u.r2_key, { httpMetadata: { contentType: "application/octet-stream" } }); |
| 45 | parts.push(await mp.uploadPart(parts.length + 1, buf)); |
| 46 | buf = new Uint8Array(PART_SIZE); |
| 47 | fill = 0; |
| 48 | }; |
| 49 | |
| 50 | if (body) { |
| 51 | const reader = body.getReader(); |
| 52 | for (;;) { |
| 53 | const { done, value } = await reader.read(); |
| 54 | if (done) break; |
| 55 | if (!value?.length) continue; |
| 56 | added += value.length; |
| 57 | // Counted as the bytes arrive: a chunked body declares no length to check up front. |
| 58 | if (limit && u.size + added > limit.bytes) { |
| 59 | // The upload can't finish, so none of it is kept: earlier chunks, parts this request |
| 60 | // already sent and the multipart upload it may have started (the caller drops the row). |
| 61 | await reader.cancel().catch(() => {}); |
| 62 | await discard(env, { ...u, multipart_id: mp?.uploadId ?? u.multipart_id }); |
| 63 | throw limit.error(); |
| 64 | } |
| 65 | hasher?.update(value); |
| 66 | for (let off = 0; off < value.length; ) { |
| 67 | const n = Math.min(PART_SIZE - fill, value.length - off); |
| 68 | buf.set(value.subarray(off, off + n), fill); |
| 69 | fill += n; |
| 70 | off += n; |
| 71 | if (fill === PART_SIZE) await flush(); |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | const digest = final && hasher ? `sha256:${hasher.hex()}` : null; |
| 76 | |
| 77 | const hadTail = !!u.multipart_id && u.tail > 0; |
| 78 | if (final) { |
| 79 | if (mp) { |
| 80 | if (fill > 0 || !parts.length) parts.push(await mp.uploadPart(parts.length + 1, buf.subarray(0, fill))); |
| 81 | await mp.complete(parts); |
| 82 | if (hadTail) await bucket.delete(tailKey(u.id)); |
| 83 | } else if (added > 0 || u.size === 0) { |
| 84 | await bucket.put(u.r2_key, buf.subarray(0, fill), { httpMetadata: { contentType: "application/octet-stream" } }); |
| 85 | } |
| 86 | } else if (mp) { |
| 87 | if (fill > 0) await bucket.put(tailKey(u.id), buf.subarray(0, fill)); |
| 88 | else if (hadTail) await bucket.delete(tailKey(u.id)); |
| 89 | } else if (added > 0) { |
| 90 | await bucket.put(u.r2_key, buf.subarray(0, fill)); |
| 91 | } |
| 92 | // A small upload that grew into a multipart one: its first bytes went into part 1, so the |
| 93 | // single object written by an earlier chunk is replaced when the multipart upload completes. |
| 94 | return { ...u, multipart_id: mp?.uploadId ?? null, parts: JSON.stringify(parts), size: u.size + added, tail: mp && !final ? fill : 0, digest, hash_state: hasher && !final ? hasher.save() : null }; |
| 95 | } |
| 96 | |
| 97 | /** SHA-256 of a stored object, for uploads that arrived in several requests. */ |
| 98 | export async function digestOf(env: Env, key: string): Promise<{ digest: string; size: number } | null> { |
| 99 | const obj = await env.BUCKET.get(key); |
| 100 | if (!obj) return null; |
| 101 | const ds = new crypto.DigestStream("SHA-256"); |
| 102 | await obj.body.pipeTo(ds); |
| 103 | return { digest: `sha256:${hex(await ds.digest)}`, size: obj.size }; |
| 104 | } |
| 105 | |
| 106 | /** Throws away an upload's bytes. */ |
| 107 | export async function discard(env: Env, u: UploadRow): Promise<void> { |
| 108 | if (u.multipart_id) await env.BUCKET.resumeMultipartUpload(u.r2_key, u.multipart_id).abort().catch(() => {}); |
| 109 | await env.BUCKET.delete([u.r2_key, tailKey(u.id)]); |
| 110 | } |