From 8a6f07b9eb3de0a56e788f2e216ebf1224a57162 Mon Sep 17 00:00:00 2001
From: Antonio <acheong@student.dalat.org>
Date: Sun, 28 May 2023 20:12:52 +0800
Subject: [PATCH] customAccessToken Closes #38

---
 handlers.go   | 17 ++++++++---------
 middleware.go |  2 +-
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/handlers.go b/handlers.go
index 4b4fb8a..ed47c96 100644
--- a/handlers.go
+++ b/handlers.go
@@ -80,17 +80,16 @@ func nightmare(c *gin.Context) {
 	}
 	// Convert the chat request to a ChatGPT request
 	translated_request := chatgpt.ConvertAPIRequest(original_request)
-	// c.JSON(200, chatgpt_request)
 
-	// authHeader := c.GetHeader("Authorization")
+	authHeader := c.GetHeader("Authorization")
 	token := ACCESS_TOKENS.GetToken()
-	// if authHeader != "" {
-	// 	customAccessToken := strings.Replace(authHeader, "Bearer ", "", 1)
-	// 	if customAccessToken != "" {
-	// 		token = customAccessToken
-	// 		println("customAccessToken set:" + customAccessToken)
-	// 	}
-	// }
+	if authHeader != "" {
+		customAccessToken := strings.Replace(authHeader, "Bearer ", "", 1)
+		// Check if customAccessToken starts with sk-
+		if strings.HasPrefix(customAccessToken, "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6Ik1UaEVOVUpHTkVNMVFURTRNMEZCTWpkQ05UZzVNRFUxUlRVd1FVSkRNRU13UmtGRVFrRXpSZyJ9") {
+			token = customAccessToken
+		}
+	}
 
 	response, err := chatgpt.SendRequest(translated_request, token)
 	if err != nil {
diff --git a/middleware.go b/middleware.go
index 6e5328a..e87db3c 100644
--- a/middleware.go
+++ b/middleware.go
@@ -49,7 +49,7 @@ func Authorization(c *gin.Context) {
 			}
 		}
 	}
-	if len(API_KEYS) != 0 && API_KEYS[c.Request.Header.Get("Authorization")] != true {
+	if len(API_KEYS) != 0 && !API_KEYS[c.Request.Header.Get("Authorization")] {
 		c.String(401, "Unauthorized")
 		c.Abort()
 		return