From e848b44166644ca93a605db205a334891ced85e2 Mon Sep 17 00:00:00 2001
From: Frank Taylor <7483580+shagamemnon@users.noreply.github.com>
Date: Mon, 1 Mar 2021 13:19:36 -0500
Subject: [PATCH] Frank/#165 fix (#166)

* defined response status on cached responses. fixes 165

* updated mocks and added tests for range requests

* fixed range request test
---
 packages/kv-asset-handler/src/index.ts        | 20 +++++++++-
 packages/kv-asset-handler/src/mocks.ts        | 38 ++++++++++++++++---
 .../src/test/getAssetFromKV.ts                | 15 ++++++++
 3 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/packages/kv-asset-handler/src/index.ts b/packages/kv-asset-handler/src/index.ts
index 0763d81d1808..b942cab697ad 100644
--- a/packages/kv-asset-handler/src/index.ts
+++ b/packages/kv-asset-handler/src/index.ts
@@ -211,8 +211,24 @@ const getAssetFromKV = async (event: FetchEvent, options?: Partial<Options>): Pr
       }
       response = new Response(null, response)
     } else {
-      response = new Response(response.body, response)
-      response.headers.set('cf-cache-status', 'HIT')
+      headers.set('CF-Cache-Status', 'HIT')
+      // fixes #165
+      let opts = {
+        headers,
+        status: 0,
+        statusText: ''
+      }
+      if (response.status) {
+        opts.status = response.status
+        opts.statusText = response.statusText
+      } else if (headers.has('Content-Range')) {
+        opts.status = 206
+        opts.statusText = 'Partial Content'
+      } else {
+        opts.status = 200
+        opts.statusText = 'OK'
+      }
+      response = new Response(response.body, opts)
     }
 
   } else {
diff --git a/packages/kv-asset-handler/src/mocks.ts b/packages/kv-asset-handler/src/mocks.ts
index 0169ba9812ac..1e686a1b7414 100644
--- a/packages/kv-asset-handler/src/mocks.ts
+++ b/packages/kv-asset-handler/src/mocks.ts
@@ -66,17 +66,43 @@ export const mockCaches = () => {
           url: key.url,
           headers: {}
         }
+        let response
         if (key.headers.has('if-none-match')) {
-          let makeStrongEtag = key.headers.get('if-none-match').replace('W/', '')
-          Reflect.set(cacheKey.headers, 'etag', makeStrongEtag)
+          cacheKey.headers = {
+            'etag': key.headers.get('if-none-match')
+          }
+          response = cacheStore.get(JSON.stringify(cacheKey))
+        } else {
+          // if client doesn't send if-none-match, we need to iterate through these keys
+          // and just test the URL
+          const activeCacheKeys: Array<string> = Array.from(cacheStore.keys())
+          for (const cacheStoreKey of activeCacheKeys) {
+            if (JSON.parse(cacheStoreKey).url === key.url) {
+              response = cacheStore.get(cacheStoreKey)
+            }
+          }
+        }
+        if (response) {
+          // this appears overly verbose, but is necessary to document edge cache behavior
+          // The Range request header triggers the response header Content-Range ...
+          const range = key.headers.get('range')
+          if (range) {
+            response.headers.set('content-range', `bytes ${range.split('=').pop()}/${response.headers.get('content-length')}`)
+          }
+          // ... which we are using in this repository to set status 206
+          if (response.headers.has('content-range')) {
+            response.status = 206
+          } else {
+            response.status = 200
+          }
         }
-        return cacheStore.get(JSON.stringify(cacheKey))
+        return response
       },
       async put (key: any, val: Response) {
         let headers = new Headers(val.headers)
-        let url = new URL(key.url)
-        let resWithBody = new Response(val.body, { headers, status: 200 })
-        let resNoBody = new Response(null, { headers, status: 304 })
+        let body = await val.text()
+        let resp = new Response(body, { headers })
+        headers.set('content-length', (body.length).toString())
         let cacheKey: CacheKey = {
           url: key.url,
           headers: {
diff --git a/packages/kv-asset-handler/src/test/getAssetFromKV.ts b/packages/kv-asset-handler/src/test/getAssetFromKV.ts
index 17958beea449..cacc097b7db2 100644
--- a/packages/kv-asset-handler/src/test/getAssetFromKV.ts
+++ b/packages/kv-asset-handler/src/test/getAssetFromKV.ts
@@ -439,4 +439,19 @@ test('getAssetFromKV if-none-match not sent but resource in cache, should return
   }
 })
 
+test('getAssetFromKV if range request submitted and resource in cache, request fulfilled', async t => {
+  const resourceKey = 'cache.html'
+  const event1 = getEvent(new Request(`https://blah.com/${resourceKey}`))
+  const event2 = getEvent(new Request(`https://blah.com/${resourceKey}`, { headers: { 'range': 'bytes=0-10'}}))
+  const res1 = await getAssetFromKV(event1, { cacheControl: { edgeTTL: 720 } })
+  await res1
+  await sleep(2)
+  const res2 = await getAssetFromKV(event2)
+  if (res2.headers.has('content-range')) {
+    t.is(res2.status, 206)
+  } else {
+    t.fail('Response was undefined')
+  }
+})
+
 test.todo('getAssetFromKV when body not empty, should invoke .cancel()')
\ No newline at end of file