API key authentication

.mockery.yaml

@@ -26,5 +26,4 @@ packages:
       inpackage: True
       dir: "{{.InterfaceDir}}"
     interfaces:
-      Dispatcher:
-# testonly: True
+      Dispatcher:

Dockerfile

@@ -1,4 +1,4 @@
-FROM --platform=$BUILDPLATFORM golang:1.22 as builder
+FROM --platform=$BUILDPLATFORM golang:1.22 AS builder
 
 WORKDIR /app
 

README.md

@@ -48,6 +48,10 @@ GLOBAL OPTIONS:
    --log-level value   set the log level. Options: debug, info, warn, error, panic, fatal. [$LOG_LEVEL]
    --version           print the version
 
+   auth
+
+   --auth-key value, -k value  the authentication key to use for incoming requests. [$AUTH_KEY]
+
    function
 
    --arg value, -a value [ --arg value, -a value ]  additional arguments for to the worker process. [$FUNCTION_ARGS]

app/app.go

@@ -1,13 +1,14 @@
 package app
 
 import (
+	"github.com/urfave/cli/v2"
+	"go.uber.org/fx"
+
 	"github.com/lambda-feedback/shimmy/config"
 	"github.com/lambda-feedback/shimmy/internal/shell"
 	"github.com/lambda-feedback/shimmy/runtime"
 	"github.com/lambda-feedback/shimmy/util/conf"
 	"github.com/lambda-feedback/shimmy/util/logging"
-	"github.com/urfave/cli/v2"
-	"go.uber.org/fx"
 )
 
 func New(ctx *cli.Context) (*shell.Shell, error) {

app/lambda/handler.go

@@ -7,9 +7,10 @@ import (
 
 	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/awslabs/aws-lambda-go-api-proxy/httpadapter"
-	"github.com/lambda-feedback/shimmy/internal/server"
 	"go.uber.org/fx"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/server"
 )
 
 // LambdaHandlerParams represents the parameters required for

app/lambda/module.go

@@ -1,9 +1,10 @@
 package lambda
 
 import (
+	"go.uber.org/fx"
+
 	"github.com/lambda-feedback/shimmy/handler"
 	"github.com/lambda-feedback/shimmy/util/logging"
-	"go.uber.org/fx"
 )
 
 func Module(config Config) fx.Option {

app/standalone/module.go

@@ -1,10 +1,11 @@
 package standalone
 
 import (
+	"go.uber.org/fx"
+
 	"github.com/lambda-feedback/shimmy/handler"
 	"github.com/lambda-feedback/shimmy/internal/server"
 	"github.com/lambda-feedback/shimmy/util/logging"
-	"go.uber.org/fx"
 )
 
 func Module(config Config) fx.Option {

cmd/lambda.go

@@ -1,11 +1,12 @@
 package cmd
 
 import (
+	"github.com/urfave/cli/v2"
+
 	"github.com/lambda-feedback/shimmy/app"
 	"github.com/lambda-feedback/shimmy/app/lambda"
 	"github.com/lambda-feedback/shimmy/util/conf"
 	"github.com/lambda-feedback/shimmy/util/logging"
-	"github.com/urfave/cli/v2"
 )
 
 var (

cmd/root.go

@@ -6,11 +6,12 @@ import (
 	"os"
 	"time"
 
+	"github.com/urfave/cli/v2"
+	"go.uber.org/zap"
+
 	"github.com/lambda-feedback/shimmy/config"
 	"github.com/lambda-feedback/shimmy/util/conf"
 	"github.com/lambda-feedback/shimmy/util/logging"
-	"github.com/urfave/cli/v2"
-	"go.uber.org/zap"
 )
 
 var (
@@ -35,6 +36,13 @@ functions on arbitrary, serverless platforms.`
 				Usage:   "set the log format. Options: production, development.",
 				EnvVars: []string{"LOG_FORMAT"},
 			},
+			// auth flags
+			&cli.StringFlag{
+				Name:     "auth-key",
+				Usage:    "the secret key for the application.",
+				Category: "auth",
+				EnvVars:  []string{"AUTH_KEY"},
+			},
 			// shim flags
 			&cli.StringFlag{
 				Name:     "interface",
@@ -241,6 +249,7 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) {
 
 	// map cli flags to config fields
 	cliMap := map[string]string{
+		"auth-key":                   "auth.key",
 		"max-workers":                "runtime.max_workers",
 		"command":                    "runtime.cmd",
 		"cwd":                        "runtime.cwd",

cmd/run.go

@@ -3,8 +3,9 @@ package cmd
 import (
 	"os"
 
-	"github.com/lambda-feedback/shimmy/util/logging"
 	"github.com/urfave/cli/v2"
+
+	"github.com/lambda-feedback/shimmy/util/logging"
 )
 
 var (

cmd/serve.go

@@ -1,11 +1,12 @@
 package cmd
 
 import (
+	"github.com/urfave/cli/v2"
+
 	"github.com/lambda-feedback/shimmy/app"
 	"github.com/lambda-feedback/shimmy/app/standalone"
 	"github.com/lambda-feedback/shimmy/util/conf"
 	"github.com/lambda-feedback/shimmy/util/logging"
-	"github.com/urfave/cli/v2"
 )
 
 var (

config/config.go

@@ -8,6 +8,11 @@ const (
 	JSON MessageEncoding = "json"
 )
 
+type AuthConfig struct {
+	// Key is the secret key for the application
+	Key string `conf:"key"`
+}
+
 type Config struct {
 	// LogLevel is the log level for the application
 	LogLevel string `conf:"log_level"`
@@ -17,4 +22,7 @@ type Config struct {
 
 	// Runtime is the runtime configuration
 	Runtime runtime.Config `conf:"runtime"`
+
+	// Auth is the authentication configuration
+	Auth AuthConfig `conf:"auth"`
 }

handler/handler.go

@@ -4,27 +4,32 @@ import (
 	"io"
 	"net/http"
 
-	"github.com/lambda-feedback/shimmy/runtime"
 	"go.uber.org/fx"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/config"
+	"github.com/lambda-feedback/shimmy/runtime"
 )
 
 type CommandHandlerParams struct {
 	fx.In
 
 	Handler runtime.Handler
+	Config  config.Config
 	Log     *zap.Logger
 }
 
 func NewCommandHandler(params CommandHandlerParams) *CommandHandler {
 	return &CommandHandler{
 		handler: params.Handler,
+		config:  params.Config,
 		log:     params.Log,
 	}
 }
 
 type CommandHandler struct {
 	handler runtime.Handler
+	config  config.Config
 	log     *zap.Logger
 }
 
@@ -34,6 +39,13 @@ func (h *CommandHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		zap.String("method", r.Method),
 	)
 
+	// Check for authorization
+	if h.config.Auth.Key != "" && r.Header.Get("api-key") != h.config.Auth.Key {
+		log.Debug("unauthorized request")
+		http.Error(w, "unauthorized", http.StatusUnauthorized)
+		return
+	}
+
 	body, err := io.ReadAll(r.Body)
 	if err != nil {
 		log.Debug("failed to read body", zap.Error(err))

handler/handler_test.go

@@ -0,0 +1,102 @@
+package handler
+
+import (
+	"bytes"
+	"context"
+	"github.com/lambda-feedback/shimmy/config"
+	"github.com/lambda-feedback/shimmy/runtime"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"go.uber.org/zap"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+// --- Mock handler ---
+type MockHandler struct {
+	mock.Mock
+}
+
+func (m *MockHandler) Handle(ctx context.Context, req runtime.Request) runtime.Response {
+	args := m.Called(ctx, req)
+	return args.Get(0).(runtime.Response)
+}
+
+// --- Test ---
+func TestServeHTTP_Success(t *testing.T) {
+	mockHandler := new(MockHandler)
+
+	reqBody := []byte(`{"example": "value"}`)
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader(reqBody))
+	req.Header.Set("api-key", "secret")
+
+	w := httptest.NewRecorder()
+
+	expectedResponse := runtime.Response{
+		StatusCode: http.StatusOK,
+		Header:     http.Header{"Content-Type": []string{"application/json"}},
+		Body:       []byte(`{"ok":true}`),
+	}
+
+	mockHandler.On("Handle", mock.Anything, mock.MatchedBy(func(r runtime.Request) bool {
+		return r.Path == "/test" &&
+			r.Method == http.MethodPost &&
+			bytes.Equal(r.Body, reqBody)
+	})).Return(expectedResponse)
+
+	handler := &CommandHandler{
+		handler: mockHandler,
+		log:     zap.NewNop(), // or zaptest.NewLogger(t)
+		config: config.Config{
+			LogLevel: "debug",
+			Runtime:  runtime.Config{},
+			Auth:     config.AuthConfig{Key: "secret"},
+		},
+	}
+
+	handler.ServeHTTP(w, req)
+
+	res := w.Result()
+	defer res.Body.Close()
+
+	body, _ := io.ReadAll(res.Body)
+
+	assert.Equal(t, http.StatusOK, res.StatusCode)
+	assert.Equal(t, "application/json", res.Header.Get("Content-Type"))
+	assert.Equal(t, `{"ok":true}`, string(body))
+	mockHandler.AssertExpectations(t)
+}
+
+func TestServeHTTP_Unauthorized(t *testing.T) {
+	mockHandler := new(MockHandler)
+
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader([]byte(`{"example": "value"}`)))
+	req.Header.Set("api-key", "wrong-key") // wrong key
+
+	w := httptest.NewRecorder()
+
+	handler := &CommandHandler{
+		handler: mockHandler, // won't be called
+		log:     zap.NewNop(),
+		config: config.Config{
+			LogLevel: "debug",
+			Runtime:  runtime.Config{},
+			Auth:     config.AuthConfig{Key: "Secret"},
+		},
+	}
+
+	handler.ServeHTTP(w, req)
+
+	res := w.Result()
+	defer res.Body.Close()
+
+	body, _ := io.ReadAll(res.Body)
+
+	assert.Equal(t, http.StatusUnauthorized, res.StatusCode)
+	assert.Contains(t, string(body), "unauthorized")
+
+	// Ensure handler was not called
+	mockHandler.AssertNotCalled(t, "Handle", mock.Anything, mock.Anything)
+}

internal/execution/dispatcher.go

@@ -3,9 +3,10 @@ package execution
 import (
 	"context"
 
+	"go.uber.org/zap"
+
 	"github.com/lambda-feedback/shimmy/internal/execution/dispatcher"
 	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
-	"go.uber.org/zap"
 )
 
 type Dispatcher dispatcher.Dispatcher

internal/execution/dispatcher/dispatcher_dedicated.go

@@ -4,8 +4,9 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
 )
 
 type DedicatedDispatcher struct {

internal/execution/dispatcher/dispatcher_pooled.go

@@ -6,8 +6,9 @@ import (
 	"runtime"
 
 	"github.com/jackc/puddle/v2"
-	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
 )
 
 type PooledDispatcher struct {

internal/execution/dispatcher/dispatcher_pooled_test.go

@@ -5,11 +5,12 @@ import (
 	"testing"
 	"time"
 
-	"github.com/lambda-feedback/shimmy/internal/execution/dispatcher"
-	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/execution/dispatcher"
+	"github.com/lambda-feedback/shimmy/internal/execution/supervisor"
 )
 
 func TestPooledDispatcher_New_UsesCPUCoreFallback(t *testing.T) {

internal/execution/supervisor/adapter.go

@@ -4,8 +4,9 @@ import (
 	"context"
 	"time"
 
-	"github.com/lambda-feedback/shimmy/internal/execution/worker"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/execution/worker"
 )
 
 // AdapterWorkerFactoryFn is a type alias for a function that creates a worker

internal/execution/supervisor/adapter_file.go

@@ -12,8 +12,9 @@ import (
 	"sync"
 	"time"
 
-	"github.com/lambda-feedback/shimmy/internal/execution/worker"
 	"go.uber.org/zap"
+
+	"github.com/lambda-feedback/shimmy/internal/execution/worker"
 )
 
 // fileAdapter is an adapter that allows supervisors to use files to