From a666782384e4382f21c6a823ded4f1092ec44161 Mon Sep 17 00:00:00 2001
From: Simon Larsen <simonlarsen@oneuptime.com>
Date: Sun, 14 May 2023 19:46:22 +0100
Subject: [PATCH] fix response headers

---
 CommonServer/Utils/Response.ts | 40 +++++++++++++++++++---------------
 TestServer/API/Main.ts         | 12 +++++++++-
 TestServer/API/Settings.ts     |  9 ++++++++
 3 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/CommonServer/Utils/Response.ts b/CommonServer/Utils/Response.ts
index e6ef73030f..407273bb27 100644
--- a/CommonServer/Utils/Response.ts
+++ b/CommonServer/Utils/Response.ts
@@ -15,6 +15,7 @@ import BaseModel from 'Common/Models/BaseModel';
 import EmptyResponse from 'Common/Types/API/EmptyResponse';
 import JSONFunctions from 'Common/Types/JSONFunctions';
 import FileModel from 'Common/Models/FileModel';
+import Dictionary from 'Common/Types/Dictionary';
 
 export default class Response {
     private static logResponse(
@@ -29,20 +30,16 @@ export default class Response {
         const method: string = oneUptimeRequest.method;
         const url: URL = URL.fromString(oneUptimeRequest.url);
 
-        const header_info: string = `Response ID: ${
-            oneUptimeRequest.id
-        } -- POD NAME: ${
-            process.env['POD_NAME'] || 'NONE'
-        } -- METHOD: ${method} -- URL: ${url.toString()} -- DURATION: ${(
-            requestEndedAt.getTime() -
-            (oneUptimeRequest.requestStartedAt as Date).getTime()
-        ).toString()}ms -- STATUS: ${oneUptimeResponse.statusCode}`;
+        const header_info: string = `Response ID: ${oneUptimeRequest.id
+            } -- POD NAME: ${process.env['POD_NAME'] || 'NONE'
+            } -- METHOD: ${method} -- URL: ${url.toString()} -- DURATION: ${(
+                requestEndedAt.getTime() -
+                (oneUptimeRequest.requestStartedAt as Date).getTime()
+            ).toString()}ms -- STATUS: ${oneUptimeResponse.statusCode}`;
 
-        const body_info: string = `Response ID: ${
-            oneUptimeRequest.id
-        } -- RESPONSE BODY: ${
-            responsebody ? JSON.stringify(responsebody, null, 2) : 'EMPTY'
-        }`;
+        const body_info: string = `Response ID: ${oneUptimeRequest.id
+            } -- RESPONSE BODY: ${responsebody ? JSON.stringify(responsebody, null, 2) : 'EMPTY'
+            }`;
 
         if (oneUptimeResponse.statusCode > 299) {
             logger.error(header_info + '\n ' + body_info);
@@ -73,11 +70,18 @@ export default class Response {
         req: ExpressRequest,
         res: ExpressResponse,
         statusCode: number,
-        body: JSONObject | string
+        body: JSONObject | string,
+        headers: Dictionary<string>
     ): void {
         const oneUptimeRequest: OneUptimeRequest = req as OneUptimeRequest;
         const oneUptimeResponse: OneUptimeResponse = res as OneUptimeResponse;
 
+        if (headers) {
+            for (const key in headers) {
+                oneUptimeResponse.set(key, headers[key]?.toString() || '');
+            }
+        }
+
         oneUptimeResponse.set(
             'ExpressRequest-Id',
             oneUptimeRequest.id.toString()
@@ -158,7 +162,7 @@ export default class Response {
         res: ExpressResponse,
         list: Array<BaseModel>,
         count: PositiveNumber,
-        modelType: { new (): BaseModel }
+        modelType: { new(): BaseModel }
     ): void {
         return this.sendJsonArrayResponse(
             req,
@@ -174,15 +178,15 @@ export default class Response {
         req: ExpressRequest,
         res: ExpressResponse,
         item: BaseModel | null,
-        modelType: { new (): BaseModel }
+        modelType: { new(): BaseModel }
     ): void {
         return this.sendJsonObjectResponse(
             req,
             res,
             item
                 ? JSONFunctions.serialize(
-                      JSONFunctions.toJSONObject(item, modelType)
-                  )
+                    JSONFunctions.toJSONObject(item, modelType)
+                )
                 : {}
         );
     }
diff --git a/TestServer/API/Main.ts b/TestServer/API/Main.ts
index 78662ce83e..ffc6f84d41 100644
--- a/TestServer/API/Main.ts
+++ b/TestServer/API/Main.ts
@@ -7,6 +7,8 @@ import Express, {
 } from 'CommonServer/Utils/Express';
 import Response from 'CommonServer/Utils/Response';
 import Sleep from 'Common/Types/Sleep';
+import Typeof from 'Common/Types/Typeof';
+import { JSONValue } from 'Common/Types/JSON';
 
 const router: ExpressRouter = Express.getRouter();
 
@@ -24,6 +26,13 @@ router.get(
                 LocalCache.getNumber('TestServer', 'responseTime') || 0;
             const responseBody: string | undefined =
                 LocalCache.getString('TestServer', 'responseBody') || '';
+            let responseHeaders: JSONValue | undefined =
+                LocalCache.getJSON('TestServer', 'responseHeaders') || {};
+
+            
+            if(responseHeaders && typeof responseHeaders === Typeof.String) {
+                responseHeaders = JSON.parse(responseHeaders.toString());
+            }
 
             if (responseTime > 0) {
                 await Sleep.sleep(responseTime);
@@ -35,7 +44,8 @@ router.get(
                 req,
                 res,
                 responseCode,
-                responseBody
+                responseBody,
+                responseHeaders ? responseHeaders as any : {}
             );
         } catch (err) {
             return next(err);
diff --git a/TestServer/API/Settings.ts b/TestServer/API/Settings.ts
index 70150ca7e2..612a1ed1b8 100644
--- a/TestServer/API/Settings.ts
+++ b/TestServer/API/Settings.ts
@@ -33,6 +33,15 @@ router.post(
                 | string
                 | undefined;
 
+            const responseHeaders: JSONObject | undefined = data['responseHeaders'] as
+                | JSONObject
+                | undefined;
+
+            LocalCache.setJSON(
+                'TestServer',
+                'responseHeaders',
+                responseHeaders || ''
+            );
             LocalCache.setString(
                 'TestServer',
                 'responseType',