optional authentication

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

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)
 }

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()
+}