Added Check API KEYs file to API.go (#1381)

Added API KEYs file

Signed-off-by: lunamidori5 <118759930+lunamidori5@users.noreply.github.com>
This commit is contained in:
lunamidori5 2023-12-04 19:06:45 -08:00 committed by GitHub
parent 67966b623c
commit 563c5b7ea0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -3,6 +3,8 @@ package api
import ( import (
"errors" "errors"
"fmt" "fmt"
"encoding/json"
"io/ioutil"
"strings" "strings"
config "github.com/go-skynet/LocalAI/api/config" config "github.com/go-skynet/LocalAI/api/config"
@ -144,30 +146,48 @@ func App(opts ...options.AppOption) (*fiber.App, error) {
// Auth middleware checking if API key is valid. If no API key is set, no auth is required. // Auth middleware checking if API key is valid. If no API key is set, no auth is required.
auth := func(c *fiber.Ctx) error { auth := func(c *fiber.Ctx) error {
if len(options.ApiKeys) > 0 { if len(options.ApiKeys) == 0 {
authHeader := c.Get("Authorization") return c.Next()
if authHeader == "" { }
return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"message": "Authorization header missing"})
} // Check for api_keys.json file
authHeaderParts := strings.Split(authHeader, " ") fileContent, err := ioutil.ReadFile("api_keys.json")
if len(authHeaderParts) != 2 || authHeaderParts[0] != "Bearer" { if err == nil {
return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"message": "Invalid Authorization header format"}) // Parse JSON content from the file
var fileKeys []string
err := json.Unmarshal(fileContent, &fileKeys)
if err != nil {
return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{"message": "Error parsing api_keys.json"})
} }
apiKey := authHeaderParts[1] // Add file keys to options.ApiKeys
validApiKey := false options.ApiKeys = append(options.ApiKeys, fileKeys...)
for _, key := range options.ApiKeys { }
if apiKey == key {
validApiKey = true authHeader := c.Get("Authorization")
} if authHeader == "" {
} return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"message": "Authorization header missing"})
if !validApiKey { }
return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"message": "Invalid API key"}) authHeaderParts := strings.Split(authHeader, " ")
if len(authHeaderParts) != 2 || authHeaderParts[0] != "Bearer" {
return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"message": "Invalid Authorization header format"})
}
apiKey := authHeaderParts[1]
validApiKey := false
for _, key := range options.ApiKeys {
if apiKey == key {
validApiKey = true
} }
} }
if !validApiKey {
return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"message": "Invalid API key"})
}
return c.Next() return c.Next()
} }
if options.CORS { if options.CORS {
var c func(ctx *fiber.Ctx) error var c func(ctx *fiber.Ctx) error
if options.CORSAllowOrigins == "" { if options.CORSAllowOrigins == "" {