From 8f748d637d991725b1515b15fe54bf7097733a28 Mon Sep 17 00:00:00 2001
From: Kalinka George <xitowzys@gmail.com>
Date: Wed, 26 Apr 2023 16:17:06 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=A6=20chore:=20Change=20CORS?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 core/rest_api/api.py | 57 +++++++++++++++++++++++++++++++++++---------
 1 file changed, 46 insertions(+), 11 deletions(-)

diff --git a/core/rest_api/api.py b/core/rest_api/api.py
index a20b069..45712b9 100644
--- a/core/rest_api/api.py
+++ b/core/rest_api/api.py
@@ -1,5 +1,6 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from loguru import logger
+from starlette.responses import JSONResponse
 from uvicorn_loguru_integration import run_uvicorn_loguru
 import uvicorn
 
@@ -8,6 +9,32 @@
 from fastapi.middleware.cors import CORSMiddleware
 
 
+async def custom_cors_middleware(request: Request, call_next):
+    response = None
+    allowed_port = 9000
+    origin = request.headers.get("origin")
+
+    if origin and origin.endswith(f":{allowed_port}"):
+        response = await call_next(request)
+
+        response.headers["Access-Control-Allow-Origin"] = origin
+        response.headers["Access-Control-Allow-Credentials"] = "true"
+        response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, DELETE, OPTIONS"
+        response.headers["Access-Control-Allow-Headers"] = "Authorization, Content-Type, Accept"
+
+    if request.method == "OPTIONS" and response:
+        response = JSONResponse(content={})
+        response.headers["Access-Control-Allow-Origin"] = origin
+        response.headers["Access-Control-Allow-Credentials"] = "true"
+        response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, DELETE, OPTIONS"
+        response.headers["Access-Control-Allow-Headers"] = "Authorization, Content-Type, Accept"
+
+    if not response:
+        response = await call_next(request)
+
+    return response
+
+
 class API:
     def __init__(self):
         self.__app: FastAPI = FastAPI(debug=True)
@@ -16,7 +43,7 @@ def __init__(self):
 
         config: fastapi_gateway_auto_generate.Config = fastapi_gateway_auto_generate.Config(
             fast_api_app=self.__app,
-            db_path=DICT_ENVS["DATABASE_PATH"] # "./database/database.db"
+            db_path=DICT_ENVS["DATABASE_PATH"]  # "./database/database.db"
         )
 
         fastapi_gateway_auto_generate.Generator(config=config)
@@ -25,15 +52,23 @@ def __init_routes(self):
         logger.success("Routers initialized")
 
     def __cors_settings(self):
-        origins = ["*"]
-
-        self.__app.add_middleware(
-            CORSMiddleware,
-            allow_origins=origins,
-            allow_credentials=True,
-            allow_methods=["*"],
-            allow_headers=["*"],
-        )
+        # def custom_origin_check(origin: str) -> bool:
+        #     # Check if the request comes from port 9000
+        #     return origin.endswith(':9000')
+
+        # origins = ["*"]
+        #
+        # self.__app.add_middleware(
+        #     CORSMiddleware,
+        #     allow_origin_callback=origins,
+        #     allow_credentials=True,
+        #     allow_methods=["*"],
+        #     allow_headers=["*"],
+        #     expose_headers=["*"],
+        #
+        # )
+
+        self.__app.middleware("http")(custom_cors_middleware)
 
         logger.success("Cors settings are installed")