Define types and structure for generic middleware creation (#1298)

dmjb · web-flow · commit ae59712989e5 · 2025-08-07T10:53:25.000+01:00
This is the first step towards generalizing the handling of middleware in the code. This is intended to allow new middleware to be added without altering the runner code or the runconfig schema.

This changes the runconfig to contain a list of generic MiddlewareConfig structs, and modifies the runner to create middleware instances based on this list, assuming a set of factory functions are supplied to the runner.

Right now, this list is never populated, and none of the existing middleware have been converted to use this scheme. This will be done in a future PR.
diff --git a/cmd/thv/app/proxy.go b/cmd/thv/app/proxy.go
@@ -195,7 +195,7 @@ func proxyCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 
 	// Create middlewares slice for incoming request authentication
-	var middlewares []types.Middleware
+	var middlewares []types.MiddlewareFunction
 
 	// Get OIDC configuration if enabled (for protecting the proxy endpoint)
 	var oidcConfig *auth.TokenValidatorConfig
@@ -514,7 +514,7 @@ func resolveClientSecret() (string, error) {
 }
 
 // createTokenInjectionMiddleware creates a middleware that injects the OAuth token into requests
-func createTokenInjectionMiddleware(tokenSource *oauth2.TokenSource) types.Middleware {
+func createTokenInjectionMiddleware(tokenSource *oauth2.TokenSource) types.MiddlewareFunction {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			token, err := (*tokenSource).Token()
diff --git a/docs/server/docs.go b/docs/server/docs.go
diff --git a/docs/server/swagger.json b/docs/server/swagger.json
diff --git a/docs/server/swagger.yaml b/docs/server/swagger.yaml
diff --git a/pkg/audit/config.go b/pkg/audit/config.go
@@ -8,6 +8,8 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+
+	"github.com/stacklok/toolhive/pkg/transport/types"
 )
 
 // Config represents the audit logging configuration.
@@ -103,7 +105,7 @@ func (c *Config) ShouldAuditEvent(eventType string) bool {
 }
 
 // CreateMiddleware creates an HTTP middleware from the audit configuration.
-func (c *Config) CreateMiddleware() (func(http.Handler) http.Handler, error) {
+func (c *Config) CreateMiddleware() (types.MiddlewareFunction, error) {
 	auditor, err := NewAuditor(c)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create auditor: %w", err)
diff --git a/pkg/authz/config.go b/pkg/authz/config.go
@@ -10,6 +10,8 @@ import (
 	"strings"
 
 	"sigs.k8s.io/yaml"
+
+	"github.com/stacklok/toolhive/pkg/transport/types"
 )
 
 // ConfigType represents the type of authorization configuration.
@@ -117,7 +119,7 @@ func (c *Config) Validate() error {
 }
 
 // CreateMiddleware creates an HTTP middleware from the configuration.
-func (c *Config) CreateMiddleware() (func(http.Handler) http.Handler, error) {
+func (c *Config) CreateMiddleware() (types.MiddlewareFunction, error) {
 	// Create the appropriate middleware based on the configuration type
 	switch c.Type {
 	case ConfigTypeCedarV1:
diff --git a/pkg/mcp/tool_filter.go b/pkg/mcp/tool_filter.go
@@ -32,7 +32,7 @@ var errBug = errors.New("there's a bug")
 // This middleware is designed to be used ONLY when tool filtering is enabled,
 // and expects the list of tools to be "correct" (i.e. not empty and not
 // containing nonexisting tools).
-func NewToolFilterMiddleware(filterTools []string) (types.Middleware, error) {
+func NewToolFilterMiddleware(filterTools []string) (types.MiddlewareFunction, error) {
 	if len(filterTools) == 0 {
 		return nil, fmt.Errorf("tools list for filtering is empty")
 	}
@@ -73,7 +73,7 @@ func NewToolFilterMiddleware(filterTools []string) (types.Middleware, error) {
 // This middleware is designed to be used ONLY when tool filtering is enabled,
 // and expects the list of tools to be "correct" (i.e. not empty and not
 // containing nonexisting tools).
-func NewToolCallFilterMiddleware(filterTools []string) (types.Middleware, error) {
+func NewToolCallFilterMiddleware(filterTools []string) (types.MiddlewareFunction, error) {
 	if len(filterTools) == 0 {
 		return nil, fmt.Errorf("tools list for filtering is empty")
 	}
diff --git a/pkg/runner/config.go b/pkg/runner/config.go
@@ -128,6 +128,10 @@ type RunConfig struct {
 
 	// IgnoreConfig contains configuration for ignore processing
 	IgnoreConfig *ignore.Config `json:"ignore_config,omitempty" yaml:"ignore_config,omitempty"`
+
+	// MiddlewareConfigs contains the list of middleware to apply to the transport
+	// and the configuration for each middleware.
+	MiddlewareConfigs []types.MiddlewareConfig `json:"middleware_configs,omitempty" yaml:"middleware_configs,omitempty"`
 }
 
 // WriteJSON serializes the RunConfig to JSON and writes it to the provided writer
@@ -167,6 +171,12 @@ func (c *RunConfig) WithAudit(config *audit.Config) *RunConfig {
 	return c
 }
 
+// WithMiddlewareConfig adds middleware configuration to the RunConfig
+func (c *RunConfig) WithMiddlewareConfig(middlewareConfig []types.MiddlewareConfig) *RunConfig {
+	c.MiddlewareConfigs = middlewareConfig
+	return c
+}
+
 // WithTransport parses and sets the transport type
 func (c *RunConfig) WithTransport(t string) (*RunConfig, error) {
 	transportType, err := types.ParseTransportType(t)
diff --git a/pkg/runner/config_builder.go b/pkg/runner/config_builder.go
@@ -58,6 +58,12 @@ func (b *RunConfigBuilder) WithName(name string) *RunConfigBuilder {
 	return b
 }
 
+// WithMiddlewareConfig sets the middleware configuration
+func (b *RunConfigBuilder) WithMiddlewareConfig(middlewareConfig []types.MiddlewareConfig) *RunConfigBuilder {
+	b.config.MiddlewareConfigs = middlewareConfig
+	return b
+}
+
 // WithCmdArgs sets the command arguments
 func (b *RunConfigBuilder) WithCmdArgs(args []string) *RunConfigBuilder {
 	b.config.CmdArgs = args
diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
@@ -29,6 +29,9 @@ type Runner struct {
 
 	// telemetryProvider is the OpenTelemetry provider for cleanup
 	telemetryProvider *telemetry.Provider
+
+	// supportedMiddleware is a map of supported middleware types to their factory functions.
+	supportedMiddleware map[string]types.MiddlewareFactory
 }
 
 // NewRunner creates a new Runner with the provided configuration
@@ -53,6 +56,29 @@ func (r *Runner) Run(ctx context.Context) error {
 		Debug:      r.Config.Debug,
 	}
 
+	// Create middleware from the MiddlewareConfigs instances in the RunConfig.
+	for _, middlewareConfig := range r.Config.MiddlewareConfigs {
+		// First, get the correct factory function for the middleware type.
+		factory, ok := r.supportedMiddleware[middlewareConfig.Type]
+		if !ok {
+			return fmt.Errorf("unsupported middleware type: %s", middlewareConfig.Type)
+		}
+
+		// Create the middleware instance using the factory function.
+		middleware, err := factory(&middlewareConfig)
+		if err != nil {
+			return fmt.Errorf("failed to create middleware of type %s: %v", middlewareConfig.Type, err)
+		}
+
+		// Ensure middleware is cleaned up on shutdown.
+		defer func() {
+			if err := middleware.Close(); err != nil {
+				logger.Warnf("Failed to close middleware of type %s: %v", middlewareConfig.Type, err)
+			}
+		}()
+		transportConfig.Middlewares = append(transportConfig.Middlewares, middleware.Handler())
+	}
+
 	if len(r.Config.ToolsFilter) > 0 {
 		toolsFilterMiddleware, err := mcp.NewToolFilterMiddleware(r.Config.ToolsFilter)
 		if err != nil {
diff --git a/pkg/telemetry/config.go b/pkg/telemetry/config.go
@@ -290,7 +290,7 @@ func createMetricExporter(ctx context.Context, config Config) (sdkmetric.Exporte
 // Middleware returns an HTTP middleware that instruments requests with OpenTelemetry.
 // serverName is the name of the MCP server (e.g., "github", "fetch")
 // transport is the backend transport type ("stdio" or "sse")
-func (p *Provider) Middleware(serverName, transport string) types.Middleware {
+func (p *Provider) Middleware(serverName, transport string) types.MiddlewareFunction {
 	return NewHTTPMiddleware(p.config, p.tracerProvider, p.meterProvider, serverName, transport)
 }
 
diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go
@@ -48,7 +48,7 @@ func NewHTTPMiddleware(
 	tracerProvider trace.TracerProvider,
 	meterProvider metric.MeterProvider,
 	serverName, transport string,
-) types.Middleware {
+) types.MiddlewareFunction {
 	meter := meterProvider.Meter(instrumentationName)
 
 	// Initialize metrics
diff --git a/pkg/transport/http.go b/pkg/transport/http.go
@@ -33,7 +33,7 @@ type HTTPTransport struct {
 	containerName     string
 	deployer          rt.Deployer
 	debug             bool
-	middlewares       []types.Middleware
+	middlewares       []types.MiddlewareFunction
 	prometheusHandler http.Handler
 
 	// Mutex for protecting shared state
@@ -60,7 +60,7 @@ func NewHTTPTransport(
 	debug bool,
 	targetHost string,
 	prometheusHandler http.Handler,
-	middlewares ...types.Middleware,
+	middlewares ...types.MiddlewareFunction,
 ) *HTTPTransport {
 	if host == "" {
 		host = LocalhostIPv4
diff --git a/pkg/transport/proxy/httpsse/http_proxy.go b/pkg/transport/proxy/httpsse/http_proxy.go
@@ -52,7 +52,7 @@ type HTTPSSEProxy struct {
 	host          string
 	port          int
 	containerName string
-	middlewares   []types.Middleware
+	middlewares   []types.MiddlewareFunction
 
 	// HTTP server
 	server     *http.Server
@@ -78,7 +78,7 @@ type HTTPSSEProxy struct {
 
 // NewHTTPSSEProxy creates a new HTTP SSE proxy for transports.
 func NewHTTPSSEProxy(
-	host string, port int, containerName string, prometheusHandler http.Handler, middlewares ...types.Middleware,
+	host string, port int, containerName string, prometheusHandler http.Handler, middlewares ...types.MiddlewareFunction,
 ) *HTTPSSEProxy {
 	proxy := &HTTPSSEProxy{
 		middlewares:       middlewares,
@@ -100,7 +100,7 @@ func NewHTTPSSEProxy(
 }
 
 // applyMiddlewares applies a chain of middlewares to a handler
-func applyMiddlewares(handler http.Handler, middlewares ...types.Middleware) http.Handler {
+func applyMiddlewares(handler http.Handler, middlewares ...types.MiddlewareFunction) http.Handler {
 	// Apply middleware chain in reverse order (last middleware is applied first)
 	for i := len(middlewares) - 1; i >= 0; i-- {
 		handler = middlewares[i](handler)
diff --git a/pkg/transport/proxy/streamable/streamable_proxy.go b/pkg/transport/proxy/streamable/streamable_proxy.go
@@ -28,7 +28,7 @@ type HTTPProxy struct {
 	containerName     string
 	shutdownCh        chan struct{}
 	prometheusHandler http.Handler
-	middlewares       []types.Middleware
+	middlewares       []types.MiddlewareFunction
 
 	// Message channel for sending JSON-RPC to the container
 	messageCh chan jsonrpc2.Message
@@ -44,7 +44,7 @@ func NewHTTPProxy(
 	port int,
 	containerName string,
 	prometheusHandler http.Handler,
-	middlewares ...types.Middleware,
+	middlewares ...types.MiddlewareFunction,
 ) *HTTPProxy {
 	return &HTTPProxy{
 		host:              host,
diff --git a/pkg/transport/proxy/transparent/transparent_proxy.go b/pkg/transport/proxy/transparent/transparent_proxy.go
@@ -44,7 +44,7 @@ type TransparentProxy struct {
 	server *http.Server
 
 	// Middleware chain
-	middlewares []types.Middleware
+	middlewares []types.MiddlewareFunction
 
 	// Mutex for protecting shared state
 	mutex sync.Mutex
@@ -76,7 +76,7 @@ func NewTransparentProxy(
 	targetURI string,
 	prometheusHandler http.Handler,
 	enableHealthCheck bool,
-	middlewares ...types.Middleware,
+	middlewares ...types.MiddlewareFunction,
 ) *TransparentProxy {
 	proxy := &TransparentProxy{
 		host:              host,
diff --git a/pkg/transport/stdio.go b/pkg/transport/stdio.go
@@ -32,7 +32,7 @@ type StdioTransport struct {
 	containerName     string
 	deployer          rt.Deployer
 	debug             bool
-	middlewares       []types.Middleware
+	middlewares       []types.MiddlewareFunction
 	prometheusHandler http.Handler
 
 	// Mutex for protecting shared state
@@ -61,7 +61,7 @@ func NewStdioTransport(
 	deployer rt.Deployer,
 	debug bool,
 	prometheusHandler http.Handler,
-	middlewares ...types.Middleware,
+	middlewares ...types.MiddlewareFunction,
 ) *StdioTransport {
 	return &StdioTransport{
 		host:              host,
diff --git a/pkg/transport/types/transport.go b/pkg/transport/types/transport.go
@@ -4,6 +4,7 @@ package types
 
 import (
 	"context"
+	"encoding/json"
 	"net/http"
 
 	"golang.org/x/exp/jsonrpc2"
@@ -14,16 +15,52 @@ import (
 	"github.com/stacklok/toolhive/pkg/transport/errors"
 )
 
-// Middleware is a function that wraps an http.Handler with additional functionality.
-type Middleware func(http.Handler) http.Handler
+// MiddlewareFunction is a function that wraps an http.Handler with additional functionality.
+type MiddlewareFunction func(http.Handler) http.Handler
+
+// Middleware defines a middleware interceptor and a close method.
+type Middleware interface {
+	// Handler returns the middleware function used by the proxy.
+	Handler() MiddlewareFunction
+	// Close cleans up any resources used by the middleware.
+	Close() error
+}
+
+// MiddlewareConfig represents the configuration for a middleware.
+// This is stored in the run config, and is used to instantiate the middleware.
+type MiddlewareConfig struct {
+	// Type is a string representing the middleware type.
+	Type string `json:"type"`
+	// Parameters is a JSON object containing the middleware parameters.
+	// It is stored as a raw message to allow flexible parameter types.
+	Parameters json.RawMessage `json:"parameters" swaggertype:"object"`
+}
+
+// NewMiddlewareConfig creates a new MiddlewareConfig with the given type and parameters.
+// The parameters are marshaled to JSON and stored in the Parameters field.
+func NewMiddlewareConfig(middlewareType string, parameters any) (*MiddlewareConfig, error) {
+	// Marshal the parameters to JSON
+	paramsJSON, err := json.Marshal(parameters)
+	if err != nil {
+		return nil, err
+	}
+
+	return &MiddlewareConfig{
+		Type:       middlewareType,
+		Parameters: paramsJSON,
+	}, nil
+}
+
+// MiddlewareFactory is a function that creates a Middleware instance based on the provided configuration.
+type MiddlewareFactory func(config *MiddlewareConfig) (Middleware, error)
 
 // Transport defines the interface for MCP transport implementations.
 // It provides methods for handling communication between the client and server.
 type Transport interface {
 	// Mode returns the transport mode.
 	Mode() TransportType
 
-	// Port returns the port used by the transport.
+	// ProxyPort returns the port used by the transport.
 	ProxyPort() int
 
 	// Setup prepares the transport for use.
@@ -132,7 +169,7 @@ type Config struct {
 
 	// Middlewares is a list of middleware functions to apply to the transport.
 	// These are applied in order, with the first middleware being the outermost wrapper.
-	Middlewares []Middleware
+	Middlewares []MiddlewareFunction
 
 	// PrometheusHandler is an optional HTTP handler for Prometheus metrics endpoint.
 	// If provided, it will be exposed at /metrics on the transport's HTTP server.

Original file line number	Diff line number	Diff line change
`@@ -195,7 +195,7 @@ func proxyCmdFunc(cmd *cobra.Command, args []string) error {`
`195`	`195`	`}`
`196`	`196`
`197`	`197`	`// Create middlewares slice for incoming request authentication`
`198`		`- var middlewares []types.Middleware`
	`198`	`+ var middlewares []types.MiddlewareFunction`
`199`	`199`
`200`	`200`	`// Get OIDC configuration if enabled (for protecting the proxy endpoint)`
`201`	`201`	`var oidcConfig *auth.TokenValidatorConfig`
`@@ -514,7 +514,7 @@ func resolveClientSecret() (string, error) {`
`514`	`514`	`}`
`515`	`515`
`516`	`516`	`// createTokenInjectionMiddleware creates a middleware that injects the OAuth token into requests`
`517`		`-func createTokenInjectionMiddleware(tokenSource *oauth2.TokenSource) types.Middleware {`
	`517`	`+func createTokenInjectionMiddleware(tokenSource *oauth2.TokenSource) types.MiddlewareFunction {`
`518`	`518`	`return func(next http.Handler) http.Handler {`
`519`	`519`	`return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {`
`520`	`520`	`token, err := (*tokenSource).Token()`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,8 @@ import (`
`10`	`10`	`"strings"`
`11`	`11`
`12`	`12`	`"sigs.k8s.io/yaml"`
	`13`	`+`
	`14`	`+ "github.com/stacklok/toolhive/pkg/transport/types"`
`13`	`15`	`)`
`14`	`16`
`15`	`17`	`// ConfigType represents the type of authorization configuration.`
`@@ -117,7 +119,7 @@ func (c *Config) Validate() error {`
`117`	`119`	`}`
`118`	`120`
`119`	`121`	`// CreateMiddleware creates an HTTP middleware from the configuration.`
`120`		`-func (c *Config) CreateMiddleware() (func(http.Handler) http.Handler, error) {`
	`122`	`+func (c *Config) CreateMiddleware() (types.MiddlewareFunction, error) {`
`121`	`123`	`// Create the appropriate middleware based on the configuration type`
`122`	`124`	`switch c.Type {`
`123`	`125`	`case ConfigTypeCedarV1:`