diff --git a/handlers.go b/handlers.go index 1ef20cb..0718a17 100644 --- a/handlers.go +++ b/handlers.go @@ -31,6 +31,22 @@ func passwordHandler(c *gin.Context) { c.String(200, "password updated") } +func puidHandler(c *gin.Context) { + // Get the password from the request (json) and update the password + type puid_struct struct { + PUID string `json:"puid"` + } + var puid puid_struct + err := c.BindJSON(&puid) + if err != nil { + c.String(400, "puid not provided") + return + } + // Set environment variable + os.Setenv("PUID", puid.PUID) + c.String(200, "puid updated") +} + func tokensHandler(c *gin.Context) { // Get the request_tokens from the request (json) and update the request_tokens var request_tokens []string diff --git a/main.go b/main.go index 15d6c68..e6bdbbb 100644 --- a/main.go +++ b/main.go @@ -69,9 +69,10 @@ func main() { /// Admin routes admin_routes.PATCH("/password", passwordHandler) - admin_routes.PATCH("/tokens", adminCheck, tokensHandler) + admin_routes.PATCH("/tokens", tokensHandler) + admin_routes.PATCH("/puid", puidHandler) /// Public routes router.OPTIONS("/v1/chat/completions", optionsHandler) - router.POST("/v1/chat/completions", nightmare) + router.POST("/v1/chat/completions", Authorization, nightmare) endless.ListenAndServe(HOST+":"+PORT, router) } diff --git a/middleware.go b/middleware.go index 3884788..c52c76d 100644 --- a/middleware.go +++ b/middleware.go @@ -31,3 +31,13 @@ func cors(c *gin.Context) { c.Header("Access-Control-Allow-Headers", "*") c.Next() } + +func Authorization(c *gin.Context) { + api_password := os.Getenv("API_PASSWORD") + if api_password != "" && c.Request.Header.Get("Authorization") != api_password { + c.String(401, "Unauthorized") + c.Abort() + return + } + c.Next() +}