guard

Documentation

Overview

Package guard provides security context management and guard registry for the MCP Gateway.

This package is responsible for managing security labels (DIFC - Decentralized Information Flow Control) and storing/retrieving agent identifiers in request contexts.

Relationship with internal/auth:

- internal/auth: Primary authentication logic (header parsing, validation) - internal/guard: Security context management (agent ID tracking, guard registry)

For authentication-related operations, always use the internal/auth package directly.

Example:

// Extract agent ID from auth header and store in context
agentID := auth.ExtractAgentID(authHeader)
ctx = guard.SetAgentIDInContext(ctx, agentID)

// Retrieve agent ID from context
agentID := guard.GetAgentIDFromContext(ctx) // Returns "default" if not found

Index

• func BuildLabelAgentPayload(policy interface{}, trustedBots []string) interface{}
• func GetAgentIDFromContext(ctx context.Context) string
• func GetRegisteredGuardTypes() []string
• func RegisterGuardType(name string, factory GuardFactory)
• func SetAgentIDInContext(ctx context.Context, agentID string) context.Context
• func SetRequestStateInContext(ctx context.Context, state RequestState) context.Context
• type AgentLabelsPayload
• type BackendCaller
• type ContextKey
• type Guard
  • func CreateGuard(name string) (Guard, error)
• type GuardFactory
• type LabelAgentResult
• type NoopGuard
  • func NewNoopGuard() *NoopGuard
  • func (g *NoopGuard) LabelAgent(ctx context.Context, policy interface{}, backend BackendCaller, ...) (*LabelAgentResult, error)
  • func (g *NoopGuard) LabelResource(ctx context.Context, toolName string, args interface{}, backend BackendCaller, ...) (*difc.LabeledResource, difc.OperationType, error)
  • func (g *NoopGuard) LabelResponse(ctx context.Context, toolName string, result interface{}, ...) (difc.LabeledData, error)
  • func (g *NoopGuard) Name() string
• type Registry
  • func NewRegistry() *Registry
  • func (r *Registry) Get(serverID string) Guard
  • func (r *Registry) GetGuardInfo() map[string]string
  • func (r *Registry) Has(serverID string) bool
  • func (r *Registry) HasNonNoopGuard() bool
  • func (r *Registry) List() []string
  • func (r *Registry) Register(serverID string, guard Guard)
  • func (r *Registry) Remove(serverID string)
• type RequestState
  • func GetRequestStateFromContext(ctx context.Context) RequestState
• type WasmGuard
  • func NewWasmGuard(ctx context.Context, name string, wasmPath string, backend BackendCaller) (*WasmGuard, error)
  • func NewWasmGuardFromBytes(ctx context.Context, name string, wasmBytes []byte, backend BackendCaller) (*WasmGuard, error)
  • func NewWasmGuardWithOptions(ctx context.Context, name string, wasmBytes []byte, backend BackendCaller, ...) (*WasmGuard, error)
  • func (g *WasmGuard) Close(ctx context.Context) error
  • func (g *WasmGuard) LabelAgent(ctx context.Context, policy interface{}, backend BackendCaller, ...) (*LabelAgentResult, error)
  • func (g *WasmGuard) LabelResource(ctx context.Context, toolName string, args interface{}, backend BackendCaller, ...) (*difc.LabeledResource, difc.OperationType, error)
  • func (g *WasmGuard) LabelResponse(ctx context.Context, toolName string, result interface{}, ...) (difc.LabeledData, error)
  • func (g *WasmGuard) Name() string
• type WasmGuardOptions
• type WriteSinkGuard
  • func NewWriteSinkGuard(accept []string) *WriteSinkGuard
  • func (g *WriteSinkGuard) LabelAgent(_ context.Context, _ interface{}, _ BackendCaller, _ *difc.Capabilities) (*LabelAgentResult, error)
  • func (g *WriteSinkGuard) LabelResource(_ context.Context, toolName string, _ interface{}, _ BackendCaller, ...) (*difc.LabeledResource, difc.OperationType, error)
  • func (g *WriteSinkGuard) LabelResponse(_ context.Context, _ string, _ interface{}, _ BackendCaller, ...) (difc.LabeledData, error)
  • func (g *WriteSinkGuard) Name() string

Functions

func BuildLabelAgentPayload

func BuildLabelAgentPayload(policy interface{}, trustedBots []string) interface{}

BuildLabelAgentPayload constructs the label_agent input payload from the given guard policy and an optional list of additional trusted bot usernames. The trusted bots are merged with the guard's built-in list and cannot remove any built-in entries. If trustedBots is nil or empty, the returned payload contains only the allow-only policy.

func GetAgentIDFromContext

func GetAgentIDFromContext(ctx context.Context) string

GetAgentIDFromContext extracts the agent ID from the context Returns "default" if not found

func GetRegisteredGuardTypes

func GetRegisteredGuardTypes() []string

GetRegisteredGuardTypes returns all registered guard type names

func RegisterGuardType

func RegisterGuardType(name string, factory GuardFactory)

RegisterGuardType registers a guard type with a factory function This allows dynamic guard creation by name

func SetAgentIDInContext

func SetAgentIDInContext(ctx context.Context, agentID string) context.Context

SetAgentIDInContext sets the agent ID in the context

func SetRequestStateInContext

func SetRequestStateInContext(ctx context.Context, state RequestState) context.Context

SetRequestStateInContext stores guard request state in context

Types

type AgentLabelsPayload

type AgentLabelsPayload struct {
	Secrecy   []string `json:"secrecy"`
	Integrity []string `json:"integrity"`
}

AgentLabelsPayload holds effective secrecy/integrity labels for the session.

type BackendCaller

type BackendCaller interface {
	// CallTool makes a read-only call to the backend MCP server
	// This is used by guards to gather metadata for labeling
	CallTool(ctx context.Context, toolName string, args interface{}) (interface{}, error)
}

BackendCaller provides a way for guards to make read-only calls to the backend to gather information needed for labeling (e.g., fetching issue author)

type ContextKey

type ContextKey string

ContextKey is used for storing values in context

const (
	// AgentIDContextKey stores the agent ID in the request context
	AgentIDContextKey ContextKey = "difc-agent-id"

	// RequestStateContextKey stores guard-specific request state
	RequestStateContextKey ContextKey = "difc-request-state"
)

type Guard

type Guard interface {
	// Name returns the identifier for this guard (e.g., "github", "noop")
	Name() string

	// LabelAgent initializes guard policy and returns effective agent/session state
	// for the current session.
	// Returns:
	//   - result: effective labels, mode, and normalized policy
	//   - error: any validation/initialization error
	LabelAgent(ctx context.Context, policy interface{}, backend BackendCaller, caps *difc.Capabilities) (*LabelAgentResult, error)

	// LabelResource determines the resource being accessed and its labels
	// This may call the backend (via BackendCaller) to gather metadata needed for labeling
	// Returns:
	//   - resource: The labeled resource (simple or nested structure for fine-grained filtering)
	//   - operation: The type of operation (Read, Write, or ReadWrite)
	//   - error: Any error that occurred during labeling
	LabelResource(ctx context.Context, toolName string, args interface{}, backend BackendCaller, caps *difc.Capabilities) (*difc.LabeledResource, difc.OperationType, error)

	// LabelResponse labels the response data after a successful backend call
	// This is used for fine-grained filtering of collections
	// Returns:
	//   - labeledData: The response data with per-item labels (if applicable)
	//   - error: Any error that occurred during labeling
	// If the guard returns nil for labeledData, the reference monitor will use the
	// resource labels from LabelResource for the entire response
	LabelResponse(ctx context.Context, toolName string, result interface{}, backend BackendCaller, caps *difc.Capabilities) (difc.LabeledData, error)
}

Guard handles DIFC labeling for a specific MCP server Guards ONLY label resources - they do NOT make access control decisions The Reference Monitor (in the server) uses guard-provided labels to enforce DIFC policies

func CreateGuard

func CreateGuard(name string) (Guard, error)

CreateGuard creates a guard instance by name using registered factories

type GuardFactory

type GuardFactory func() (Guard, error)

GuardFactory is a function that creates a guard instance

type LabelAgentResult

type LabelAgentResult struct {
	Agent            AgentLabelsPayload     `json:"agent"`
	DIFCMode         string                 `json:"difc_mode"`
	NormalizedPolicy map[string]interface{} `json:"normalized_policy,omitempty"`
}

LabelAgentResult describes the effective policy/session state returned by a guard.

type NoopGuard

type NoopGuard struct{}

NoopGuard is the default guard that performs no DIFC labeling It allows all operations by returning empty labels (no restrictions)

func NewNoopGuard

func NewNoopGuard() *NoopGuard

NewNoopGuard creates a new noop guard

func (*NoopGuard) LabelAgent

func (g *NoopGuard) LabelAgent(ctx context.Context, policy interface{}, backend BackendCaller, caps *difc.Capabilities) (*LabelAgentResult, error)

LabelAgent initializes noop guard session state.

func (*NoopGuard) LabelResource

func (g *NoopGuard) LabelResource(ctx context.Context, toolName string, args interface{}, backend BackendCaller, caps *difc.Capabilities) (*difc.LabeledResource, difc.OperationType, error)

LabelResource returns an empty resource with no label requirements Treats all operations as read-write (most conservative assumption)

func (*NoopGuard) LabelResponse

func (g *NoopGuard) LabelResponse(ctx context.Context, toolName string, result interface{}, backend BackendCaller, caps *difc.Capabilities) (difc.LabeledData, error)

LabelResponse returns nil, indicating no fine-grained labeling The reference monitor will use the resource labels for the entire response

func (*NoopGuard) Name

func (g *NoopGuard) Name() string

Name returns the identifier for this guard

type Registry

type Registry struct {
	// contains filtered or unexported fields
}

Registry manages guard instances for different MCP servers

func NewRegistry

func NewRegistry() *Registry

NewRegistry creates a new guard registry

func (*Registry) Get

func (r *Registry) Get(serverID string) Guard

Get retrieves the guard for a server, or returns a noop guard if not found

func (*Registry) GetGuardInfo

func (r *Registry) GetGuardInfo() map[string]string

GetGuardInfo returns information about all registered guards

func (*Registry) Has

func (r *Registry) Has(serverID string) bool

Has checks if a guard is registered for a server

func (*Registry) HasNonNoopGuard

func (r *Registry) HasNonNoopGuard() bool

HasNonNoopGuard returns true if any registered guard is not a noop guard

func (*Registry) List

func (r *Registry) List() []string

List returns all registered server IDs

func (*Registry) Register

func (r *Registry) Register(serverID string, guard Guard)

Register registers a guard for a specific server

func (*Registry) Remove

func (r *Registry) Remove(serverID string)

Remove removes a guard registration

type RequestState

type RequestState interface{}

RequestState represents any state that the guard needs to pass from request to response This is useful when the guard needs to carry information from LabelResource to LabelResponse

func GetRequestStateFromContext

func GetRequestStateFromContext(ctx context.Context) RequestState

GetRequestStateFromContext retrieves guard request state from context

type WasmGuard

type WasmGuard struct {
	// contains filtered or unexported fields
}

WasmGuard implements Guard interface by executing a WASM module in-process The WASM module runs sandboxed within the gateway using wazero runtime Guards cannot make direct network calls - they receive a BackendCaller interface via host functions

Thread Safety: WASM modules are single-threaded, so all calls to a guard instance are serialized using a mutex. Concurrent requests will queue and execute one at a time.

func NewWasmGuard

func NewWasmGuard(ctx context.Context, name string, wasmPath string, backend BackendCaller) (*WasmGuard, error)

NewWasmGuard creates a new WASM guard from a WASM binary file

func NewWasmGuardFromBytes

func NewWasmGuardFromBytes(ctx context.Context, name string, wasmBytes []byte, backend BackendCaller) (*WasmGuard, error)

NewWasmGuardFromBytes creates a new WASM guard from WASM binary bytes This is useful when loading guards from URLs or other sources

func NewWasmGuardWithOptions

func NewWasmGuardWithOptions(ctx context.Context, name string, wasmBytes []byte, backend BackendCaller, opts *WasmGuardOptions) (*WasmGuard, error)

NewWasmGuardWithOptions creates a new WASM guard from WASM binary bytes with custom options Options can be nil to use defaults (stdout/stderr go to os.Stdout/os.Stderr)

func (*WasmGuard) Close

func (g *WasmGuard) Close(ctx context.Context) error

Close releases WASM runtime resources

func (*WasmGuard) LabelAgent

func (g *WasmGuard) LabelAgent(ctx context.Context, policy interface{}, backend BackendCaller, caps *difc.Capabilities) (*LabelAgentResult, error)

LabelAgent calls the WASM module's label_agent function.

func (*WasmGuard) LabelResource

func (g *WasmGuard) LabelResource(ctx context.Context, toolName string, args interface{}, backend BackendCaller, caps *difc.Capabilities) (*difc.LabeledResource, difc.OperationType, error)

LabelResource calls the WASM module's label_resource function

func (*WasmGuard) LabelResponse

func (g *WasmGuard) LabelResponse(ctx context.Context, toolName string, result interface{}, backend BackendCaller, caps *difc.Capabilities) (difc.LabeledData, error)

LabelResponse calls the WASM module's label_response function

func (*WasmGuard) Name

func (g *WasmGuard) Name() string

Name returns the identifier for this guard

type WasmGuardOptions

type WasmGuardOptions struct {
	// Stdout is the writer for WASM stdout output. Defaults to os.Stdout if nil.
	Stdout io.Writer
	// Stderr is the writer for WASM stderr output. Defaults to os.Stderr if nil.
	Stderr io.Writer
}

WasmGuardOptions configures optional settings for WASM guard creation

type WriteSinkGuard

type WriteSinkGuard struct {
	// contains filtered or unexported fields
}

WriteSinkGuard is a guard for write-only output channels (e.g., safe-outputs).

When DIFC is enabled, an agent that reads from a guarded server (like GitHub) acquires secrecy and integrity tags. Writing to an unguarded output server would fail because the DIFC evaluator sees a label mismatch: the agent has tags that the resource (with empty labels from a noop guard) does not.

The write-sink guard fixes this by:

• Returning empty labels from LabelAgent (does not contribute agent labels)
• Setting resource secrecy to the configured accept patterns
• Classifying all operations as OperationWrite

This ensures writes succeed because:

• Integrity: resource requires no tags (empty), agent has all zero required → OK
• Secrecy: resource secrecy includes the agent's secrecy patterns → agentSecrecy ⊆ resourceSecrecy → OK

Write-sink is required for ALL output servers when DIFC guards are enabled, including when repos="all" or repos="public". Without it, the noop guard assigns OperationRead + empty labels, causing integrity violations when the agent has integrity tags from other guards.

Configuration examples:

// For repos="all" or repos="public" (agent has no secrecy):
"guard-policies": {
  "write-sink": {
    "accept": ["*"]
  }
}

// For scoped repos (agent has secrecy tags):
"guard-policies": {
  "write-sink": {
    "accept": ["private:github/gh-aw*"]
  }
}

func NewWriteSinkGuard

func NewWriteSinkGuard(accept []string) *WriteSinkGuard

NewWriteSinkGuard creates a new write-sink guard with the specified accept patterns. Each pattern becomes a secrecy tag on the resource, allowing agents with matching secrecy to write to this sink.

func (*WriteSinkGuard) LabelAgent

func (g *WriteSinkGuard) LabelAgent(_ context.Context, _ interface{}, _ BackendCaller, _ *difc.Capabilities) (*LabelAgentResult, error)

LabelAgent returns empty labels. The write-sink does not contribute agent labels — those are set by the primary guard (e.g., the GitHub WASM guard).

func (*WriteSinkGuard) LabelResource

func (g *WriteSinkGuard) LabelResource(_ context.Context, toolName string, _ interface{}, _ BackendCaller, _ *difc.Capabilities) (*difc.LabeledResource, difc.OperationType, error)

LabelResource sets the resource's secrecy to the configured accept patterns and classifies the operation as a write.

For writes the DIFC evaluator checks:

• agentSecrecy ⊆ resource.Secrecy (no secret information leak)
• resource.Integrity ⊆ agentIntegrity (agent is trusted enough)

By setting the resource secrecy to the accept patterns from config, agents whose secrecy tags are a subset of the accept set can write successfully. By leaving the resource integrity empty, the second check also passes because the agent has all zero of the (empty) required integrity tags.

func (*WriteSinkGuard) LabelResponse

func (g *WriteSinkGuard) LabelResponse(_ context.Context, _ string, _ interface{}, _ BackendCaller, _ *difc.Capabilities) (difc.LabeledData, error)

LabelResponse returns nil; the write-sink does not perform fine-grained response labeling since all operations are writes (responses are confirmations).

func (*WriteSinkGuard) Name

func (g *WriteSinkGuard) Name() string

Name returns the identifier for this guard