add key auth for admin endpoints

.gitignore

@@ -1,2 +1,3 @@
 release_notes.md
-target
+target
+.DS_STORE

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.8.0 - 2024-01-17
+### Added
+- Added key authentication for admin endpoints
+
 ## 1.7.6 - 2024-01-17
 ### Fixed
 - Changed code to string in OpenAI error response

README.md

@@ -145,10 +145,18 @@ docker pull luyuanxin1995/bricksllm:1.4.0
 > | `REDIS_WRITE_TIME_OUT`         | optional | Timeout for Redis write operations | `500ms`
 > | `IN_MEMORY_DB_UPDATE_INTERVAL`         | optional | The interval BricksLLM API gateway polls Postgresql DB for latest key configurations | `1s`
 > | `STATS_PROVIDER`         | optional | This value can only be datadog. Required for integration with Datadog.  |
-> | `PROXY_TIMEOUT`         | optional | This value can only be datadog. Required for integration with Datadog.  |
+> | `PROXY_TIMEOUT`         | optional | Timeout for proxy HTTP requests. |
+> | `ADMIN_PASS`         | optional | Simple password authentication for admin endpoints.  |
 
 ## Configuration Endpoints
 The configuration server runs on Port `8001`.
+
+##### Headers
+> | name   |  type      | data type      | description                                          |
+> |--------|------------|----------------|------------------------------------------------------|
+> | `X-API-KEY` |  optional  | `string`         | Key authentication header.
+
+
 <details>
   <summary>Get keys: <code>GET</code> <code><b>/api/key-management/keys</b></code></summary>
 

cmd/bricksllm/main.go

@@ -182,7 +182,7 @@ func main() {
 	cpm := manager.NewCustomProvidersManager(store, cpMemStore)
 	rm := manager.NewRouteManager(store, store, rMemStore, psMemStore)
 
-	as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm)
+	as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm, cfg.AdminPass)
 	if err != nil {
 		log.Sugar().Fatalf("error creating admin http server: %v", err)
 	}

cmd/tool/main.go

@@ -181,7 +181,7 @@ func main() {
 	cpm := manager.NewCustomProvidersManager(store, cpMemStore)
 	rm := manager.NewRouteManager(store, store, rMemStore, psMemStore)
 
-	as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm)
+	as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm, cfg.AdminPass)
 	if err != nil {
 		log.Sugar().Fatalf("error creating admin http server: %v", err)
 	}

internal/config/config.go

@@ -24,6 +24,7 @@ type Config struct {
 	InMemoryDbUpdateInterval time.Duration `env:"IN_MEMORY_DB_UPDATE_INTERVAL" envDefault:"5s"`
 	OpenAiKey                string        `env:"OPENAI_API_KEY"`
 	StatsProvider            string        `env:"STATS_PROVIDER"`
+	AdminPass                string        `env:"ADMIN_PASS"`
 	ProxyTimeout             time.Duration `env:"PROXY_TIMEOUT" envDefault:"180s"`
 }
 

internal/server/web/admin/admin.go

@@ -55,11 +55,11 @@ type AdminServer struct {
 	m      KeyManager
 }
 
-func NewAdminServer(log *zap.Logger, mode string, m KeyManager, krm KeyReportingManager, psm ProviderSettingsManager, cpm CustomProvidersManager, rm RouteManager) (*AdminServer, error) {
+func NewAdminServer(log *zap.Logger, mode string, m KeyManager, krm KeyReportingManager, psm ProviderSettingsManager, cpm CustomProvidersManager, rm RouteManager, adminPass string) (*AdminServer, error) {
 	router := gin.New()
 
 	prod := mode == "production"
-	router.Use(getAdminLoggerMiddleware(log, "admin", prod))
+	router.Use(getAdminLoggerMiddleware(log, "admin", prod, adminPass))
 
 	router.GET("/api/health", getGetHealthCheckHandler())
 

internal/server/web/admin/middleware.go

@@ -8,8 +8,14 @@ import (
 	"go.uber.org/zap"
 )
 
-func getAdminLoggerMiddleware(log *zap.Logger, prefix string, prod bool) gin.HandlerFunc {
+func getAdminLoggerMiddleware(log *zap.Logger, prefix string, prod bool, adminPass string) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		if len(adminPass) != 0 && c.Request.Header.Get("X-API-KEY") != adminPass {
+			c.Status(200)
+			c.Abort()
+			return
+		}
+
 		c.Set(correlationId, util.NewUuid())
 		start := time.Now()
 		c.Next()