From 8623fb74d9f4e4fad231da7d5dc1ec5b0a141323 Mon Sep 17 00:00:00 2001
From: Joe Ma <rikkaneko23@gmail.com>
Date: Wed, 31 Jan 2024 03:48:28 +0800
Subject: [PATCH] Fix x-paste-read-limit handling when using Paste API v2

Fix raw/download option being ignored

Signed-off-by: Joe Ma <rikkaneko23@gmail.com>
---
 .gitignore     |  1 +
 src/index.ts   | 40 +++++++++++++++++++++++-----------------
 src/types.d.ts |  7 ++++++-
 src/utils.ts   |  5 ++++-
 wrangler.toml  |  5 ++++-
 5 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9452459..714bdea 100644
--- a/.gitignore
+++ b/.gitignore
@@ -145,6 +145,7 @@ web_modules/
 # dotenv environment variables file
 .env
 .env.test
+.dev.vars
 
 # parcel-bundler cache (https://parceljs.org/)
 .cache
diff --git a/src/index.ts b/src/index.ts
index 5a5b3b2..25908b9 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -17,13 +17,11 @@
  */
 
 import { AwsClient } from 'aws4fetch';
-import { customAlphabet } from 'nanoid';
 import { sha256 } from 'js-sha256';
 import { Router, error } from 'itty-router';
-import { ERequest } from './types';
-import { Env, PasteIndexEntry } from './types';
+import { ERequest, Env, PasteIndexEntry } from './types';
 import { serve_static } from './proxy';
-import { check_password_rules, get_paste_info, get_basic_auth } from './utils';
+import { check_password_rules, get_paste_info, get_basic_auth, gen_id } from './utils';
 import { UUID_LENGTH, PASTE_WEB_URL, SERVICE_URL } from './constant';
 
 const gen_id = customAlphabet('1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz', UUID_LENGTH);
@@ -140,14 +138,17 @@ router.post('/', async (request, env, ctx) => {
     paste_type = headers.get('x-paste-type') || undefined;
     need_qrcode = headers.get('x-paste-qr') === '1';
     reply_json = headers.get('x-json') === '1';
-    const count = headers.get('x-paste-read-limit') || '';
-    const n = parseInt(count);
-    if (isNaN(n) || n <= 0) {
-      return new Response('x-paste-read-limit must be a positive integer.\n', {
-        status: 422,
-      });
+    const count = headers.get('x-paste-read-limit') || undefined;
+    if (count) {
+      const n = parseInt(count);
+      if (isNaN(n) || n <= 0) {
+        return new Response('x-paste-read-limit must be a positive integer.\n', {
+          status: 422,
+        });
+      }
+      read_limit = n;
     }
-    read_limit = n;
+
     buffer = await request.arrayBuffer();
   }
 
@@ -370,12 +371,8 @@ router.get('/:uuid/:option?', async (request, env, ctx) => {
     // Let the browser guess the content
     else res.headers.delete('content-type');
 
-    // Handle option
-    if (option === 'raw') res.headers.delete('content-type');
-    else if (option === 'download')
-      res.headers.set('content-disposition', `attachment; filename="${encodeURIComponent(descriptor.title ?? uuid)}"`);
     // Link redirection
-    else if (descriptor.type === 'link' || option === 'link') {
+    if (descriptor.type === 'link') {
       const content = await res.clone().arrayBuffer();
       try {
         const href = new TextDecoder().decode(content);
@@ -400,6 +397,10 @@ router.get('/:uuid/:option?', async (request, env, ctx) => {
     // res.body cannot be read twice
     // Do not block when writing to cache
     if (res.ok) ctx.waitUntil(cache.put(req_key, res.clone()));
+    // Handle option
+    if (option === 'raw') res.headers.delete('content-type');
+    else if (option === 'download')
+      res.headers.set('content-disposition', `attachment; filename="${encodeURIComponent(descriptor.title ?? uuid)}"`);
     return res;
   }
 
@@ -408,7 +409,12 @@ router.get('/:uuid/:option?', async (request, env, ctx) => {
   if (res.status == 304) return res;
   let { readable, writable } = new TransformStream();
   res.body?.pipeTo(writable);
-  return new Response(readable, res);
+  const nres = new Response(readable, res);
+  // Handle option
+  if (option === 'raw') nres.headers.delete('content-type');
+  else if (option === 'download')
+    nres.headers.set('content-disposition', `attachment; filename="${encodeURIComponent(descriptor.title ?? uuid)}"`);
+  return nres;
 });
 
 // Update paste metadata
diff --git a/src/types.d.ts b/src/types.d.ts
index aeca567..a1c099a 100644
--- a/src/types.d.ts
+++ b/src/types.d.ts
@@ -14,6 +14,8 @@ export interface PasteIndexEntry {
   editable?: boolean; // Default: False (unsupported)
   read_count_remain?: number;
   type?: string;
+  upload_completed?: boolean;
+  sha256_hash?: string;
 }
 
 export interface Env {
@@ -21,5 +23,8 @@ export interface Env {
   QRCODE: ServiceWorkerGlobalScope;
   AWS_ACCESS_KEY_ID: string;
   AWS_SECRET_ACCESS_KEY: string;
+  LARGE_AWS_ACCESS_KEY_ID?: string;
+  LARGE_AWS_SECRET_ACCESS_KEY?: string;
   ENDPOINT: string;
-}
+  LARGE_ENDPOINT: string;
+}
\ No newline at end of file
diff --git a/src/utils.ts b/src/utils.ts
index dab1431..6ffb2c1 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -17,9 +17,12 @@
  */
 
 import dedent from 'dedent-js';
-import { SERVICE_URL } from './constant';
+import { customAlphabet } from 'nanoid';
+import { SERVICE_URL, UUID_LENGTH } from './constant';
 import { PasteIndexEntry, Env } from './types';
 
+export const gen_id = customAlphabet('1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz', UUID_LENGTH);
+
 export async function get_paste_info(
   uuid: string,
   descriptor: PasteIndexEntry,
diff --git a/wrangler.toml b/wrangler.toml
index 66cbe2d..f1f2cf9 100644
--- a/wrangler.toml
+++ b/wrangler.toml
@@ -12,4 +12,7 @@ services = [
 # [secret]
 # AWS_ACCESS_KEY_ID
 # AWS_SECRET_ACCESS_KEY
-# ENDPOINT
\ No newline at end of file
+# ENDPOINT
+# LARGE_AWS_ACCESS_KEY_ID
+# LARGE_AWS_SECRET_ACCESS_KEY
+# LARGE_ENDPOINT
\ No newline at end of file