feat[r2]: add basic r2 support

Author: SharzyL

src/handlers/handleDelete.ts

@@ -12,7 +12,7 @@ export async function handleDelete(request: Request, env: Env, _: ExecutionConte
     if (passwd !== metadata.passwd) {
       throw new WorkerError(403, `incorrect password for paste '${nameFromPath}`)
     } else {
-      await deletePaste(env, nameFromPath)
+      await deletePaste(env, nameFromPath, metadata)
       return new Response("the paste will be deleted in seconds")
     }
   }

src/handlers/handleRead.ts

@@ -9,6 +9,23 @@ import { parsePath } from "../shared.js"
 
 type Headers = { [name: string]: string }
 
+async function decodeMaybeStream(content: ArrayBuffer | ReadableStream): Promise<string> {
+  if (content instanceof ArrayBuffer) {
+    return decode(content)
+  } else {
+    const reader = content.pipeThrough(new TextDecoderStream()).getReader()
+    let result = ""
+    while (true) {
+      const { done, value } = await reader.read()
+      if (done) {
+        break
+      }
+      result += value
+    }
+    return result
+  }
+}
+
 function staticPageCacheHeader(env: Env): Headers {
   const age = env.CACHE_STATIC_PAGE_AGE
   return age ? { "Cache-Control": `public, max-age=${age}` } : {}
@@ -76,8 +93,8 @@ async function handleStaticPages(request: Request, env: Env, _: ExecutionContext
   return null
 }
 
-export async function handleGet(request: Request, env: Env, _: ExecutionContext): Promise<Response> {
-  const staticPageResp = await handleStaticPages(request, env, _)
+export async function handleGet(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
+  const staticPageResp = await handleStaticPages(request, env, ctx)
   if (staticPageResp !== null) {
     return staticPageResp
   }
@@ -88,7 +105,7 @@ export async function handleGet(request: Request, env: Env, _: ExecutionContext)
 
   const disp = url.searchParams.has("a") ? "attachment" : "inline"
 
-  const item = await getPaste(env, nameFromPath)
+  const item = await getPaste(env, nameFromPath, ctx)
 
   // when paste is not found
   if (item === null) {
@@ -120,7 +137,7 @@ export async function handleGet(request: Request, env: Env, _: ExecutionContext)
 
   // handle URL redirection
   if (role === "u") {
-    const redirectURL = decode(item.paste)
+    const redirectURL = await decodeMaybeStream(item.paste)
     if (isLegalUrl(redirectURL)) {
       return Response.redirect(redirectURL)
     } else {
@@ -130,7 +147,7 @@ export async function handleGet(request: Request, env: Env, _: ExecutionContext)
 
   // handle article (render as markdown)
   if (role === "a") {
-    const md = makeMarkdown(decode(item.paste))
+    const md = makeMarkdown(await decodeMaybeStream(item.paste))
     return new Response(md, {
       headers: {
         "Content-Type": `text/html;charset=UTF-8`,
@@ -143,7 +160,7 @@ export async function handleGet(request: Request, env: Env, _: ExecutionContext)
   // handle language highlight
   const lang = url.searchParams.get("lang")
   if (lang) {
-    return new Response(makeHighlight(decode(item.paste), lang), {
+    return new Response(makeHighlight(await decodeMaybeStream(item.paste), lang), {
       headers: {
         "Content-Type": `text/html;charset=UTF-8`,
         ...pasteCacheHeader(env),

src/handlers/handleWrite.ts

@@ -3,7 +3,6 @@ import { decode, genRandStr, isLegalUrl, WorkerError } from "../common.js"
 import { createPaste, getPasteMetadata, pasteNameAvailable, updatePaste } from "../storage/storage.js"
 import {
   DEFAULT_PASSWD_LEN,
-  MAX_LEN,
   NAME_REGEX,
   PASTE_NAME_LEN,
   PasteResponse,
@@ -13,6 +12,7 @@ import {
   parsePath,
   MIN_PASSWD_LEN,
   MAX_PASSWD_LEN,
+  parseSize,
 } from "../shared.js"
 
 function suggestUrl(short: string, baseUrl: string, filename?: string, contentAsString?: string) {
@@ -93,7 +93,7 @@ export async function handlePostOrPut(
   // check if paste content is legal
   if (content === undefined) {
     throw new WorkerError(400, "cannot find content in formdata")
-  } else if (contentLength > MAX_LEN) {
+  } else if (contentLength > parseSize(env.R2_MAX_ALLOWED)!) {
     throw new WorkerError(413, "payload too large")
   }
 

src/shared.ts

@@ -16,7 +16,21 @@ export const DEFAULT_PASSWD_LEN = 24
 export const MAX_PASSWD_LEN = 128
 export const MIN_PASSWD_LEN = 8
 export const PASSWD_SEP = ":"
-export const MAX_LEN = 25 * 1024 * 1024
+
+export function parseSize(sizeStr: string): number | null {
+  sizeStr = sizeStr.trim()
+  const EXPIRE_REGEX = /^[\d.]+\s*[KMG]?$/
+  if (!EXPIRE_REGEX.test(sizeStr)) {
+    return null
+  }
+
+  let sizeBytes = parseFloat(sizeStr)
+  const lastChar = sizeStr[sizeStr.length - 1]
+  if (lastChar === "K") sizeBytes *= 1024
+  else if (lastChar === "M") sizeBytes *= 1024 * 1024
+  else if (lastChar === "G") sizeBytes *= 1024 * 1024 * 1024
+  return sizeBytes
+}
 
 export function parseExpiration(expirationStr: string): number | null {
   expirationStr = expirationStr.trim()

src/storage/storage.ts

@@ -1,27 +1,79 @@
 import { dateToUnix, WorkerError } from "../common.js"
+import { parseSize } from "../shared.js"
 
+export type PasteLocation = "KV" | "R2"
+
+// TODO: allow admin to upload permanent paste
+// TODO: add filename length check
 export type PasteMetadata = {
-  schemaVersion: number
+  schemaVersion: 1
+  location: PasteLocation // new field on V1
   passwd: string
 
   lastModifiedAtUnix: number
   createdAtUnix: number
-  // TODO: allow admin to upload permanent paste
   willExpireAtUnix: number
 
   accessCounter: number // a counter representing how frequent it is accessed, to administration usage
+  sizeBytes: number
+  filename?: string
+}
+
+type PasteMetadataInStorage = {
+  schemaVersion: number
+  location?: PasteLocation
+  passwd: string
+
+  lastModifiedAtUnix: number
+  createdAtUnix: number
+  willExpireAtUnix: number
+
+  accessCounter?: number
   sizeBytes?: number
-  // TODO: add filename length check
   filename?: string
 }
 
+function migratePasteMetadata(original: PasteMetadataInStorage): PasteMetadata {
+  return {
+    schemaVersion: 1,
+    location: original.location || "KV",
+    passwd: original.passwd,
+
+    lastModifiedAtUnix: original.lastModifiedAtUnix,
+    createdAtUnix: original.createdAtUnix,
+    willExpireAtUnix: original.willExpireAtUnix,
+
+    accessCounter: original.accessCounter || 0,
+    sizeBytes: original.sizeBytes || 0,
+    filename: original.filename,
+  }
+}
+
 export type PasteWithMetadata = {
-  paste: ArrayBuffer
+  paste: ArrayBuffer | ReadableStream
   metadata: PasteMetadata
 }
 
-export async function getPaste(env: Env, short: string): Promise<PasteWithMetadata | null> {
-  const item = await env.PB.getWithMetadata<PasteMetadata>(short, {
+async function updateAccessCounter(env: Env, short: string, value: ArrayBuffer, metadata: PasteMetadata) {
+  // update counter with probability 1%
+  if (Math.random() < 0.01) {
+    metadata.accessCounter += 1
+    try {
+      await env.PB.put(short, value, {
+        metadata: metadata,
+        expiration: metadata.willExpireAtUnix,
+      })
+    } catch (e) {
+      // ignore rate limit message
+      if (!(e as Error).message.includes("KV PUT failed: 429 Too Many Requests")) {
+        throw e
+      }
+    }
+  }
+}
+
+export async function getPaste(env: Env, short: string, ctx: ExecutionContext): Promise<PasteWithMetadata | null> {
+  const item = await env.PB.getWithMetadata<PasteMetadataInStorage>(short, {
     type: "arrayBuffer",
   })
 
@@ -30,34 +82,38 @@ export async function getPaste(env: Env, short: string): Promise<PasteWithMetada
   } else if (item.metadata === null) {
     throw new WorkerError(500, `paste of name '${short}' has no metadata`)
   } else {
-    if (item.metadata.willExpireAtUnix < new Date().getTime() / 1000) {
+    const metadata = migratePasteMetadata(item.metadata)
+    const expired = metadata.willExpireAtUnix < new Date().getTime() / 1000
+
+    ctx.waitUntil(
+      (async () => {
+        if (expired) {
+          await deletePaste(env, short, metadata)
+          return null
+        }
+        await updateAccessCounter(env, short, item.value!, metadata)
+      })(),
+    )
+
+    if (expired) {
       return null
     }
 
-    // update counter with probability 1%
-    // TODO: use waitUntil API
-    if (Math.random() < 0.01) {
-      item.metadata.accessCounter += 1
-      try {
-        await env.PB.put(short, item.value, {
-          metadata: item.metadata,
-          expiration: item.metadata.willExpireAtUnix,
-        })
-      } catch (e) {
-        // ignore rate limit message
-        if (!(e as Error).message.includes("KV PUT failed: 429 Too Many Requests")) {
-          throw e
-        }
+    if (metadata.location === "R2") {
+      const object = await env.R2.get(short)
+      if (object === null) {
+        throw new WorkerError(404, `cannot find R2 bucket of name '${short}'`)
       }
+      return { paste: object.body, metadata }
+    } else {
+      return { paste: item.value, metadata }
     }
-
-    return { paste: item.value, metadata: item.metadata }
   }
 }
 
 // we separate usage of getPasteMetadata and getPaste to make access metric more reliable
 export async function getPasteMetadata(env: Env, short: string): Promise<PasteMetadata | null> {
-  const item = await env.PB.getWithMetadata<PasteMetadata>(short, {
+  const item = await env.PB.getWithMetadata<PasteMetadataInStorage>(short, {
     type: "stream",
   })
 
@@ -69,7 +125,7 @@ export async function getPasteMetadata(env: Env, short: string): Promise<PasteMe
     if (item.metadata.willExpireAtUnix < new Date().getTime() / 1000) {
       return null
     }
-    return item.metadata
+    return migratePasteMetadata(item.metadata)
   }
 }
 
@@ -87,22 +143,29 @@ export async function updatePaste(
   },
 ) {
   const expirationUnix = dateToUnix(options.now) + options.expirationSeconds
-  const putOptions: KVNamespacePutOptions = {
-    metadata: {
-      schemaVersion: 0,
-      filename: options.filename || originalMetadata.filename,
-      passwd: options.passwd,
-
-      lastModifiedAtUnix: dateToUnix(options.now),
-      createdAtUnix: originalMetadata.createdAtUnix,
-      willExpireAtUnix: expirationUnix,
-      accessCounter: originalMetadata.accessCounter,
-      sizeBytes: options.contentLength,
-    },
-    expiration: expirationUnix,
+  // since CF does not allow expiration shorter than 60s, extend the expiration to 70s
+  const expirationUnixSpecified = dateToUnix(options.now) + Math.max(options.expirationSeconds, 70)
+
+  if (originalMetadata.location === "R2") {
+    await env.R2.put(pasteName, content)
+  }
+  const metadata: PasteMetadata = {
+    schemaVersion: 1,
+    location: originalMetadata.location,
+    filename: options.filename || originalMetadata.filename,
+    passwd: options.passwd,
+
+    lastModifiedAtUnix: dateToUnix(options.now),
+    createdAtUnix: originalMetadata.createdAtUnix,
+    willExpireAtUnix: expirationUnix,
+    accessCounter: originalMetadata.accessCounter,
+    sizeBytes: options.contentLength,
   }
 
-  await env.PB.put(pasteName, content, putOptions)
+  await env.PB.put(pasteName, originalMetadata.location === "R2" ? "" : content, {
+    metadata: metadata,
+    expiration: expirationUnixSpecified,
+  })
 }
 
 export async function createPaste(
@@ -118,22 +181,32 @@ export async function createPaste(
   },
 ) {
   const expirationUnix = dateToUnix(options.now) + options.expirationSeconds
-  const putOptions: KVNamespacePutOptions = {
-    metadata: {
-      schemaVersion: 0,
-      filename: options.filename,
-      passwd: options.passwd,
-
-      lastModifiedAtUnix: dateToUnix(options.now),
-      createdAtUnix: dateToUnix(options.now),
-      willExpireAtUnix: expirationUnix,
-      accessCounter: 0,
-      sizeBytes: options.contentLength,
-    },
-    expiration: expirationUnix,
+
+  // since CF does not allow expiration shorter than 60s, extend the expiration to 70s
+  const expirationUnixSpecified = dateToUnix(options.now) + Math.max(options.expirationSeconds, 70)
+
+  const location = options.contentLength > parseSize(env.R2_THRESHOLD)! ? "R2" : "KV"
+  if (location === "R2") {
+    await env.R2.put(pasteName, content)
   }
 
-  await env.PB.put(pasteName, content, putOptions)
+  const metadata: PasteMetadata = {
+    schemaVersion: 1,
+    location: location,
+    filename: options.filename,
+    passwd: options.passwd,
+
+    lastModifiedAtUnix: dateToUnix(options.now),
+    createdAtUnix: dateToUnix(options.now),
+    willExpireAtUnix: expirationUnix,
+    accessCounter: 0,
+    sizeBytes: options.contentLength,
+  }
+
+  await env.PB.put(pasteName, location === "R2" ? "" : content, {
+    metadata: metadata,
+    expiration: expirationUnixSpecified,
+  })
 }
 
 export async function pasteNameAvailable(env: Env, pasteName: string): Promise<boolean> {
@@ -147,6 +220,9 @@ export async function pasteNameAvailable(env: Env, pasteName: string): Promise<b
   }
 }
 
-export async function deletePaste(env: Env, pasteName: string): Promise<void> {
+export async function deletePaste(env: Env, pasteName: string, originalMetadata: PasteMetadata): Promise<void> {
   await env.PB.delete(pasteName)
+  if (originalMetadata.location === "R2") {
+    await env.R2.delete(pasteName)
+  }
 }