diff --git a/README.md b/README.md index 7dc710e..2fbce28 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,7 @@ Access your agents at `http://localhost:3000` | `LOCALAGENT_STATE_DIR` | Where state gets stored | | `LOCALAGENT_LOCALRAG_URL` | LocalRAG connection | | `LOCALAGENT_ENABLE_CONVERSATIONS_LOGGING` | Toggle conversation logs | +| `LOCALAGENT_API_KEYS` | A comma separated list of api keys used for authentication | ## INSTALLATION OPTIONS diff --git a/go.mod b/go.mod index f817285..d46811f 100644 --- a/go.mod +++ b/go.mod @@ -7,8 +7,8 @@ toolchain go1.22.2 require ( github.com/bwmarrin/discordgo v0.28.1 github.com/chasefleming/elem-go v0.25.0 + github.com/dave-gray101/v2keyauth v0.0.0-20240624150259-c45d584d25e2 github.com/donseba/go-htmx v1.8.0 - github.com/dslipak/pdf v0.0.2 github.com/eritikass/githubmarkdownconvertergo v0.1.10 github.com/go-telegram/bot v1.2.1 github.com/gofiber/fiber/v2 v2.52.4 @@ -19,6 +19,7 @@ require ( github.com/philippgille/chromem-go v0.5.0 github.com/sashabaranov/go-openai v1.18.3 github.com/slack-go/slack v0.12.5 + github.com/thoj/go-ircevent v0.0.0-20210723090443-73e444401d64 github.com/tmc/langchaingo v0.1.8 github.com/traefik/yaegi v0.16.1 github.com/valyala/fasthttp v1.52.0 @@ -58,7 +59,6 @@ require ( github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect github.com/ssor/bom v0.0.0-20170718123548-6386211fdfcf // indirect github.com/temoto/robotstxt v1.1.2 // indirect - github.com/thoj/go-ircevent v0.0.0-20210723090443-73e444401d64 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/tcplisten v1.0.0 // indirect go.starlark.net v0.0.0-20230302034142-4b1e35fe2254 // indirect diff --git a/go.sum b/go.sum index 9a17a56..9be66a4 100644 --- a/go.sum +++ b/go.sum @@ -23,6 +23,8 @@ github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWR github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= +github.com/dave-gray101/v2keyauth v0.0.0-20240624150259-c45d584d25e2 h1:flLYmnQFZNo04x2NPehMbf30m7Pli57xwZ0NFqR/hb0= +github.com/dave-gray101/v2keyauth v0.0.0-20240624150259-c45d584d25e2/go.mod h1:NtWqRzAp/1tw+twkW8uuBenEVVYndEAZACWU3F3xdoQ= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -30,8 +32,6 @@ github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/donseba/go-htmx v1.8.0 h1:oTx1uUsjXZZVvcZfulZvBSPtdD1jzsvZyuK91+Q8zPE= github.com/donseba/go-htmx v1.8.0/go.mod h1:8PTAYvNKf8+QYis+DpAsggKz+sa2qljtMgvdAeNBh5s= -github.com/dslipak/pdf v0.0.2 h1:djAvcM5neg9Ush+zR6QXB+VMJzR6TdnX766HPIg1JmI= -github.com/dslipak/pdf v0.0.2/go.mod h1:2L3SnkI9cQwnAS9gfPz2iUoLC0rUZwbucpbKi5R1mUo= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/eritikass/githubmarkdownconvertergo v0.1.10 h1:mL93ADvYMOeT15DcGtK9AaFFc+RcWcy6kQBC6yS/5f4= diff --git a/main.go b/main.go index c11ca00..a233089 100644 --- a/main.go +++ b/main.go @@ -4,6 +4,7 @@ import ( "log" "os" "path/filepath" + "strings" "github.com/mudler/LocalAgent/core/state" "github.com/mudler/LocalAgent/services" @@ -18,6 +19,7 @@ var timeout = os.Getenv("LOCALAGENT_TIMEOUT") var stateDir = os.Getenv("LOCALAGENT_STATE_DIR") var localRAG = os.Getenv("LOCALAGENT_LOCALRAG_URL") var withLogs = os.Getenv("LOCALAGENT_ENABLE_CONVERSATIONS_LOGGING") == "true" +var apiKeysEnv = os.Getenv("LOCALAGENT_API_KEYS") func init() { if testModel == "" { @@ -43,6 +45,11 @@ func main() { // make sure state dir exists os.MkdirAll(stateDir, 0755) + apiKeys := []string{} + if apiKeysEnv != "" { + apiKeys = strings.Split(apiKeysEnv, ",") + } + // Create the agent pool pool, err := state.NewAgentPool( testModel, @@ -62,7 +69,7 @@ func main() { } // Create the application - app := webui.NewApp(webui.WithPool(pool)) + app := webui.NewApp(webui.WithPool(pool), webui.WithApiKeys(apiKeys...)) // Start the agents if err := pool.StartAll(); err != nil { diff --git a/webui/options.go b/webui/options.go index e20f338..f163879 100644 --- a/webui/options.go +++ b/webui/options.go @@ -5,6 +5,7 @@ import "github.com/mudler/LocalAgent/core/state" type Config struct { DefaultChunkSize int Pool *state.AgentPool + ApiKeys []string } type Option func(*Config) @@ -21,6 +22,12 @@ func WithPool(pool *state.AgentPool) Option { } } +func WithApiKeys(keys ...string) Option { + return func(c *Config) { + c.ApiKeys = keys + } +} + func (c *Config) Apply(opts ...Option) { for _, opt := range opts { opt(c) diff --git a/webui/routes.go b/webui/routes.go index cff70b1..2091eec 100644 --- a/webui/routes.go +++ b/webui/routes.go @@ -1,12 +1,16 @@ package webui import ( + "crypto/subtle" "embed" + "errors" "math/rand" "net/http" + "github.com/dave-gray101/v2keyauth" fiber "github.com/gofiber/fiber/v2" "github.com/gofiber/fiber/v2/middleware/filesystem" + "github.com/gofiber/fiber/v2/middleware/keyauth" "github.com/mudler/LocalAgent/core/agent" "github.com/mudler/LocalAgent/core/sse" "github.com/mudler/LocalAgent/core/state" @@ -27,6 +31,14 @@ func (app *App) registerRoutes(pool *state.AgentPool, webapp *fiber.App) { Browse: true, })) + if len(app.config.ApiKeys) > 0 { + kaConfig, err := GetKeyAuthConfig(app.config.ApiKeys) + if err != nil || kaConfig == nil { + panic(err) + } + webapp.Use(v2keyauth.New(*kaConfig)) + } + webapp.Get("/", func(c *fiber.Ctx) error { return c.Render("views/index", fiber.Map{ "Agents": pool.List(), @@ -139,3 +151,53 @@ func Reverse[T any](original []T) (reversed []T) { return } + +func GetKeyAuthConfig(apiKeys []string) (*v2keyauth.Config, error) { + customLookup, err := v2keyauth.MultipleKeySourceLookup([]string{"header:Authorization", "header:x-api-key", "header:xi-api-key", "cookie:token"}, keyauth.ConfigDefault.AuthScheme) + if err != nil { + return nil, err + } + + return &v2keyauth.Config{ + CustomKeyLookup: customLookup, + Next: func(c *fiber.Ctx) bool { return false }, + Validator: getApiKeyValidationFunction(apiKeys), + ErrorHandler: getApiKeyErrorHandler(false, apiKeys), + AuthScheme: "Bearer", + }, nil +} + +func getApiKeyErrorHandler(opaqueErrors bool, apiKeys []string) fiber.ErrorHandler { + return func(ctx *fiber.Ctx, err error) error { + if errors.Is(err, v2keyauth.ErrMissingOrMalformedAPIKey) { + if len(apiKeys) == 0 { + return ctx.Next() // if no keys are set up, any error we get here is not an error. + } + ctx.Set("WWW-Authenticate", "Bearer") + if opaqueErrors { + return ctx.SendStatus(401) + } + return ctx.Status(401).Render("views/login", fiber.Map{}) + } + if opaqueErrors { + return ctx.SendStatus(500) + } + return err + } +} + +func getApiKeyValidationFunction(apiKeys []string) func(*fiber.Ctx, string) (bool, error) { + + return func(ctx *fiber.Ctx, apiKey string) (bool, error) { + if len(apiKeys) == 0 { + return true, nil // If no keys are setup, accept everything + } + for _, validKey := range apiKeys { + if subtle.ConstantTimeCompare([]byte(apiKey), []byte(validKey)) == 1 { + return true, nil + } + } + return false, v2keyauth.ErrMissingOrMalformedAPIKey + } + +} diff --git a/webui/views/login.html b/webui/views/login.html new file mode 100644 index 0000000..01a5f86 --- /dev/null +++ b/webui/views/login.html @@ -0,0 +1,215 @@ + + +{{template "views/partials/header" .}} + +
+Please enter your access token to continue
+Current time (UTC): {{.CurrentDate}}
+