From 00f479aad466ca6d704fe0947cb036fb974ca834 Mon Sep 17 00:00:00 2001
From: Dmitry Tyshchenko <rainydio@gmail.com>
Date: Thu, 25 Feb 2021 13:52:24 +0200
Subject: [PATCH] fix(core-api): cache status code and headers (#4329)

---
 .../unit/core-api/service-provider.test.ts    |  6 +-
 packages/core-api/src/plugins/cache.ts        | 58 ++++++++++++++-----
 2 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/__tests__/unit/core-api/service-provider.test.ts b/__tests__/unit/core-api/service-provider.test.ts
index 37db0a2207..38ad03f46c 100644
--- a/__tests__/unit/core-api/service-provider.test.ts
+++ b/__tests__/unit/core-api/service-provider.test.ts
@@ -213,9 +213,7 @@ describe("ServiceProvider", () => {
             // @ts-ignore
             defaults.customField = "dummy";
 
-            const result = (coreApiServiceProvider.configSchema() as AnySchema).validate(
-               defaults
-            );
+            const result = (coreApiServiceProvider.configSchema() as AnySchema).validate(defaults);
 
             expect(result.error).toBeUndefined();
             expect(result.value.customField).toEqual("dummy");
@@ -346,7 +344,7 @@ describe("ServiceProvider", () => {
                 expect(result.value.plugins.cache.enabled).toEqual(false);
             });
 
-            it("should return true if process.env.CORE_API_CACHE_DISABLED is defined", async () => {
+            it("should return true if process.env.CORE_API_CACHE is defined", async () => {
                 process.env.CORE_API_CACHE = "true";
 
                 jest.resetModules();
diff --git a/packages/core-api/src/plugins/cache.ts b/packages/core-api/src/plugins/cache.ts
index 9b853e052d..5d8579782d 100644
--- a/packages/core-api/src/plugins/cache.ts
+++ b/packages/core-api/src/plugins/cache.ts
@@ -2,6 +2,12 @@ import { Crypto } from "@arkecosystem/crypto";
 import Hapi from "@hapi/hapi";
 import NodeCache from "node-cache";
 
+type CachedResponse = {
+    code: number;
+    headers: Record<string, string | string[]>;
+    payload: unknown;
+};
+
 const generateCacheKey = (request: Hapi.Request): string =>
     Crypto.HashAlgorithms.sha256(
         JSON.stringify({
@@ -34,22 +40,19 @@ export = {
             type: "onPreHandler",
             async method(request: Hapi.Request, h: Hapi.ResponseToolkit) {
                 const cacheKey: string = generateCacheKey(request);
+                const cachedResponse: CachedResponse | undefined = cache.get(cacheKey);
 
-                if (cache.has(cacheKey)) {
-                    const value: { isBoom: boolean; data: Record<string, any> } | undefined = cache.get(cacheKey);
-
-                    if (value === undefined || value === null) {
-                        return h.continue;
-                    }
+                if (cachedResponse) {
+                    const newResponse = h.response(cachedResponse.payload).code(cachedResponse.code);
 
-                    if (value.isBoom) {
-                        return h.response(value.data.payload).code(value.data.statusCode).takeover();
+                    for (const [headerName, headerValue] of Object.entries(cachedResponse.headers)) {
+                        newResponse.header(headerName, headerValue);
                     }
 
-                    return h.response(value.data).code(200).takeover();
+                    return newResponse.takeover();
+                } else {
+                    return h.continue;
                 }
-
-                return h.continue;
             },
         });
 
@@ -58,13 +61,36 @@ export = {
             async method(request: Hapi.Request, h: Hapi.ResponseToolkit) {
                 const cacheKey: string = generateCacheKey(request);
 
-                if (!cache.has(cacheKey)) {
-                    cache.set(cacheKey, {
-                        isBoom: request.response.isBoom === true,
-                        data: request.response.isBoom ? request.response.output : request.response.source,
-                    });
+                let code: number;
+                let headers: Record<string, string | string[]>;
+                let payload: unknown;
+
+                if (request.response.isBoom) {
+                    code = request.response.output.statusCode;
+                    headers = request.response.output.headers;
+                    payload = request.response.output.payload;
+                } else {
+                    code = request.response.statusCode;
+                    headers = request.response.headers;
+                    payload = request.response.source;
                 }
 
+                const cachedResponse: CachedResponse = {
+                    code,
+                    headers: {},
+                    payload,
+                };
+
+                if (code >= 300 && code < 400 && "location" in headers) {
+                    cachedResponse.headers["location"] = headers["location"];
+                }
+
+                if ("content-type" in headers) {
+                    cachedResponse["content-type"] = headers["content-type"];
+                }
+
+                cache.set(cacheKey, cachedResponse);
+
                 return h.continue;
             },
         });