feat(keys): allow to set api keys to secure the instance (#39)

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
 

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

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=

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 {

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)

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

webui/views/login.html

@@ -0,0 +1,215 @@
+<!DOCTYPE html>
+<html lang="en">
+{{template "views/partials/header" .}}
+
+<body class="bg-gradient-to-br from-gray-900 to-gray-950 text-gray-200">
+<div class="flex flex-col min-h-screen">
+
+    {{template "views/partials/menu" .}}
+
+    <div class="container mx-auto px-4 py-8 flex-grow flex items-center justify-center">
+        <!-- Auth Card -->
+        <div class="max-w-md w-full bg-gray-800/90 border border-gray-700/50 rounded-xl overflow-hidden shadow-xl">
+            <div class="animation-container">
+                <div class="text-overlay">
+                <!--    <i class="fas fa-circle-nodes text-5xl text-blue-400 mb-2"></i> -->
+                </div>
+            </div>
+            
+            <div class="p-8">
+                <div class="text-center mb-6">
+                    <h2 class="text-2xl font-bold text-white">
+                        <span class="bg-clip-text text-transparent bg-gradient-to-r from-blue-400 to-indigo-400">
+                            Authorization Required
+                        </span>
+                    </h2>
+                    <p class="text-gray-400 mt-2">Please enter your access token to continue</p>
+                </div>
+                
+                <form id="login-form" class="space-y-6" onsubmit="login(); return false;">
+                    <div>
+                        <label for="token" class="block text-sm font-medium text-gray-300 mb-2">Access Token</label>
+                        <div class="relative">
+                            <div class="absolute inset-y-0 left-0 pl-3 flex items-center pointer-events-none">
+                                <i class="fas fa-key text-gray-500"></i>
+                            </div>
+                            <input 
+                                type="password" 
+                                id="token" 
+                                name="token" 
+                                placeholder="Enter your token" 
+                                class="bg-gray-700/50 border border-gray-600 text-white placeholder-gray-400 text-sm rounded-lg focus:ring-blue-500 focus:border-blue-500 block w-full pl-10 p-2.5"
+                                required
+                            />
+                        </div>
+                    </div>
+                    
+                    <div>
+                        <button 
+                            type="submit" 
+                            class="group w-full flex items-center justify-center bg-gradient-to-r from-blue-600 to-indigo-600 hover:from-blue-700 hover:to-indigo-700 text-white py-3 px-6 rounded-lg transition duration-300 ease-in-out transform hover:scale-[1.02] hover:shadow-lg font-medium"
+                        >
+                            <i class="fas fa-sign-in-alt mr-2"></i>
+                            <span>Login</span>
+                            <i class="fas fa-arrow-right opacity-0 group-hover:opacity-100 group-hover:translate-x-2 ml-2 transition-all duration-300"></i>
+                        </button>
+                    </div>
+                </form>
+                
+                <div class="mt-8 pt-6 border-t border-gray-700/50 text-center text-sm text-gray-400">
+                    <div class="flex items-center justify-center mb-2">
+                        <i class="fas fa-shield-alt mr-2 text-blue-400"></i>
+                        <span>Instance is token protected</span>
+                    </div>
+                    <p>Current time (UTC): <span id="current-time">{{.CurrentDate}}</span></p>
+                </div>
+            </div>
+        </div>
+    </div>
+
+</div>
+
+<script>
+    function login() {
+        const token = document.getElementById('token').value;
+        if (!token.trim()) {
+            // Show error with fading effect
+            const form = document.getElementById('login-form');
+            const errorMsg = document.createElement('div');
+            errorMsg.className = 'p-3 mt-4 bg-red-900/50 text-red-200 rounded-lg border border-red-700/50 text-sm flex items-center';
+            errorMsg.innerHTML = '<i class="fas fa-exclamation-circle mr-2"></i> Please enter a valid token';
+            
+            // Remove any existing error message
+            const existingError = form.querySelector('.bg-red-900/50');
+            if (existingError) form.removeChild(existingError);
+            
+            // Add new error message with animation
+            form.appendChild(errorMsg);
+            setTimeout(() => {
+                errorMsg.style.opacity = '0';
+                errorMsg.style.transition = 'opacity 0.5s ease';
+                setTimeout(() => errorMsg.remove(), 500);
+            }, 3000);
+            return;
+        }
+        
+        var date = new Date();
+        date.setTime(date.getTime() + (24*60*60*1000));
+        document.cookie = `token=${token}; expires=${date.toGMTString()}; path=/`;
+
+        // Show loading state
+        const button = document.querySelector('button[type="submit"]');
+        const originalContent = button.innerHTML;
+        button.disabled = true;
+        button.innerHTML = '<i class="fas fa-spinner fa-spin mr-2"></i> Authenticating...';
+        button.classList.add('bg-gray-600');
+        
+        // Reload after short delay to show loading state
+        setTimeout(() => {
+            window.location.reload();
+        }, 800);
+    }
+    
+    // Update current time
+    function updateCurrentTime() {
+        const timeElement = document.getElementById('current-time');
+        if (timeElement) {
+            const now = new Date();
+            const year = now.getUTCFullYear();
+            const month = String(now.getUTCMonth() + 1).padStart(2, '0');
+            const day = String(now.getUTCDate()).padStart(2, '0');
+            const hours = String(now.getUTCHours()).padStart(2, '0');
+            const minutes = String(now.getUTCMinutes()).padStart(2, '0');
+            const seconds = String(now.getUTCSeconds()).padStart(2, '0');
+            timeElement.textContent = `${year}-${month}-${day} ${hours}:${minutes}:${seconds}`;
+        }
+    }
+    
+    // Initialize current time and update it every second
+    updateCurrentTime();
+    setInterval(updateCurrentTime, 1000);
+    
+    // Add subtle particle animation to the background
+    document.addEventListener('DOMContentLoaded', function() {
+        const animContainer = document.querySelector('.animation-container');
+        if (animContainer) {
+            const canvas = document.createElement('canvas');
+            animContainer.appendChild(canvas);
+            
+            const ctx = canvas.getContext('2d');
+            canvas.width = animContainer.offsetWidth;
+            canvas.height = animContainer.offsetHeight;
+            
+            // Create particles
+            const particles = [];
+            const particleCount = 30;
+            
+            for (let i = 0; i < particleCount; i++) {
+                particles.push({
+                    x: Math.random() * canvas.width,
+                    y: Math.random() * canvas.height,
+                    radius: Math.random() * 3 + 1,
+                    color: `rgba(${Math.random() * 50 + 50}, ${Math.random() * 100 + 100}, ${Math.random() * 155 + 100}, ${Math.random() * 0.4 + 0.1})`,
+                    speedX: Math.random() * 0.5 - 0.25,
+                    speedY: Math.random() * 0.5 - 0.25
+                });
+            }
+            
+            // Animation loop
+            function animate() {
+                requestAnimationFrame(animate);
+                ctx.clearRect(0, 0, canvas.width, canvas.height);
+                
+                particles.forEach(particle => {
+                    particle.x += particle.speedX;
+                    particle.y += particle.speedY;
+                    
+                    // Bounce off edges
+                    if (particle.x < 0 || particle.x > canvas.width) {
+                        particle.speedX = -particle.speedX;
+                    }
+                    
+                    if (particle.y < 0 || particle.y > canvas.height) {
+                        particle.speedY = -particle.speedY;
+                    }
+                    
+                    // Draw particle
+                    ctx.beginPath();
+                    ctx.arc(particle.x, particle.y, particle.radius, 0, Math.PI * 2);
+                    ctx.fillStyle = particle.color;
+                    ctx.fill();
+                });
+                
+                // Connect nearby particles with lines
+                for (let i = 0; i < particles.length; i++) {
+                    for (let j = i + 1; j < particles.length; j++) {
+                        const dx = particles[i].x - particles[j].x;
+                        const dy = particles[i].y - particles[j].y;
+                        const distance = Math.sqrt(dx * dx + dy * dy);
+                        
+                        if (distance < 100) {
+                            ctx.beginPath();
+                            ctx.moveTo(particles[i].x, particles[i].y);
+                            ctx.lineTo(particles[j].x, particles[j].y);
+                            ctx.strokeStyle = `rgba(100, 150, 255, ${0.1 * (1 - distance / 100)})`;
+                            ctx.lineWidth = 1;
+                            ctx.stroke();
+                        }
+                    }
+                }
+            }
+            
+            // Start animation
+            animate();
+            
+            // Resize handling
+            window.addEventListener('resize', () => {
+                canvas.width = animContainer.offsetWidth;
+                canvas.height = animContainer.offsetHeight;
+            });
+        }
+    });
+</script>
+
+</body>
+</html>