root/golib
1
0
Fork 0
mirror of https://github.com/therootcompany/golib.git synced 2026-11-04 06:55:00 +00:00
Code Issues Projects Releases Wiki Activity

feat: restore MCP server safety features

Browse Source
This commit is contained in:
AJ ONeal 2026-09-11 17:40:19 -06:00
parent b548ca1474
commit 55ef36d93f
No known key found for this signature in database
4 changed files with 303 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
mcp/server20260728/server.go
View File

@ -7,14 +7,19 @@ package mcpserver
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"net/http"
"net/url"
"runtime/debug"
"slices"
"strings"
"sync"
"unicode"
mcptypes "github.com/therootcompany/golib/mcp/types20260728"
)
@ -46,18 +51,30 @@ type ReadResourceRequest struct {
Params mcptypes.ReadResourceParams
}
type ToolHandlerFunc func(context.Context, CallToolRequest) (mcptypes.CallToolResult, error)
type TypedToolHandler[T mcptypes.Validatable] func(context.Context, Headers, T) (mcptypes.CallToolResult, error)
type TypedToolMiddleware[T mcptypes.Validatable] func(TypedToolHandler[T]) TypedToolHandler[T]
type ResourceHandlerFunc func(context.Context, ReadResourceRequest) (mcptypes.ReadResourceResult, error)
// RegisterTypedTool decodes a tool's arguments into T before calling handler.
// T is the application Go type for the tool input; no JSON Schema validator is
// involved.
func RegisterTypedTool[T mcptypes.Validatable](s *MCPServer, definition mcptypes.Tool, handler func(context.Context, Headers, T) (mcptypes.CallToolResult, error)) {
func RegisterTypedTool[T mcptypes.Validatable](s *MCPServer, definition mcptypes.Tool, handler TypedToolHandler[T]) {
RegisterTypedToolWithMiddleware(s, definition, handler)
}
// RegisterTypedToolWithMiddleware decodes and validates arguments before
// running middleware. Middleware receives the concrete T value.
func RegisterTypedToolWithMiddleware[T mcptypes.Validatable](s *MCPServer, definition mcptypes.Tool, handler TypedToolHandler[T], middleware ...TypedToolMiddleware[T]) {
wrapped := handler
for i := len(middleware) - 1; i >= 0; i-- {
wrapped = middleware[i](wrapped)
}
s.RegisterTool(definition, func(ctx context.Context, request CallToolRequest) (mcptypes.CallToolResult, error) {
args, err := mcptypes.DecodeValidated[T](request.Params.Arguments)
if err != nil {
return mcptypes.CallToolResult{}, fmt.Errorf("decode %s arguments: %w", definition.Name, err)
}
return handler(ctx, request.Headers, args)
return wrapped(ctx, request.Headers, args)
})
}
@ -98,6 +115,9 @@ func (s *MCPServer) RegisterTool(definition mcptypes.Tool, handler ToolHandlerFu
s.toolByName[definition.Name] = registered
}
func (s *MCPServer) RegisterResource(definition mcptypes.Resource, handler ResourceHandlerFunc) {
if err := validateResourceURI(definition.URI); err != nil {
panic(fmt.Sprintf("mcp: invalid resource URI %q: %v", definition.URI, err))
}
s.mu.Lock()
defer s.mu.Unlock()
for _, resource := range s.resources {
@ -186,6 +206,18 @@ func (s *MCPServer) discover(id *mcptypes.RequestID) any {
}
return s.result(id, raw)
}
func isAppOnly(tool mcptypes.Tool) bool {
if tool.Meta == nil || tool.Meta.UI == nil || len(tool.Meta.UI.Visibility) == 0 {
return false
}
for _, visibility := range tool.Meta.UI.Visibility {
if visibility != mcptypes.VisibilityApp {
return false
}
}
return true
}
func (s *MCPServer) capabilities() mcptypes.ServerCapabilities {
s.mu.RLock()
defer s.mu.RUnlock()
@ -200,9 +232,12 @@ func (s *MCPServer) capabilities() mcptypes.ServerCapabilities {
}
func (s *MCPServer) toolsList(id *mcptypes.RequestID) any {
s.mu.RLock()
tools := make([]mcptypes.Tool, len(s.tools))
for i := range s.tools {
tools[i] = s.tools[i].definition
tools := make([]mcptypes.Tool, 0, len(s.tools))
for _, registered := range s.tools {
if isAppOnly(registered.definition) {
continue
}
tools = append(tools, registered.definition)
}
s.mu.RUnlock()
raw, err := json.Marshal(mcptypes.ToolsListResult{ResultType: "complete", Tools: tools, CacheScope: mcptypes.CachePrivate})
@ -226,7 +261,7 @@ func (s *MCPServer) toolsCall(ctx context.Context, headers Headers, id *mcptypes
if params.Name != headers.Name {
return s.error(id, mcptypes.ErrHeaderMismatch, "Mcp-Name does not match params.name")
}
value, err := registered.handler(ctx, CallToolRequest{headers, params})
value, err := s.callToolWithRecovery(ctx, registered.handler, CallToolRequest{headers, params})
if err != nil {
if errors.Is(err, mcptypes.ErrInvalidArguments) {
return s.error(id, mcptypes.ErrInvalidParams, err.Error())
@ -257,6 +292,9 @@ func (s *MCPServer) resourceRead(ctx context.Context, headers Headers, id *mcpty
if err != nil || params.URI == "" {
return s.error(id, mcptypes.ErrInvalidParams, "Invalid params for resources/read")
}
if err := validateResourceURI(params.URI); err != nil {
return s.error(id, mcptypes.ErrInvalidParams, "Invalid resource URI")
}
s.mu.RLock()
var handler ResourceHandlerFunc
for _, resource := range s.resources {
@ -273,6 +311,16 @@ func (s *MCPServer) resourceRead(ctx context.Context, headers Headers, id *mcpty
if err != nil {
return s.error(id, mcptypes.ErrInternal, "Resource read failed")
}
for _, content := range value.Contents {
if err := validateResourceURI(content.URI); err != nil {
return s.error(id, mcptypes.ErrInternal, "Resource returned invalid URI")
}
if content.Blob != "" {
if _, err := base64.StdEncoding.DecodeString(content.Blob); err != nil {
return s.error(id, mcptypes.ErrInternal, "Resource returned invalid binary data")
}
}
}
encoded, encodeErr := json.Marshal(value)
if encodeErr != nil {
return s.error(id, mcptypes.ErrInternal, "Failed to encode result")
@ -280,6 +328,28 @@ func (s *MCPServer) resourceRead(ctx context.Context, headers Headers, id *mcpty
return s.result(id, encoded)
}
func (s *MCPServer) callToolWithRecovery(ctx context.Context, handler ToolHandlerFunc, request CallToolRequest) (result mcptypes.CallToolResult, err error) {
defer func() {
if recovered := recover(); recovered != nil {
log.Printf("[mcp] panic in tool %q: %v\n%s", request.Params.Name, recovered, debug.Stack())
result = mcptypes.NewToolResultError("Tool execution failed")
err = nil
}
}()
return handler(ctx, request)
}
func validateResourceURI(raw string) error {
if raw == "" || strings.IndexFunc(raw, func(r rune) bool { return unicode.IsControl(r) || unicode.IsSpace(r) }) >= 0 {
return errors.New("URI is empty or contains invalid characters")
}
parsed, err := url.Parse(raw)
if err != nil || parsed.Scheme == "" || parsed.Fragment != "" {
return errors.New("URI must be absolute and must not contain a fragment")
}
return nil
}
func (s *MCPServer) result(id *mcptypes.RequestID, value json.RawMessage) mcptypes.Response {
return mcptypes.Response{JSONRPC: mcptypes.JSONRPCVersion, ID: id, Result: value}
}

66
mcp/server20260728/server_test.go
View File

@ -49,6 +49,51 @@ func TestServeHTTPUsesHeadersBeforeTypedDecode(t *testing.T) {
}
}
func TestTypedMiddlewareReceivesValidatedArguments(t *testing.T) {
srv := &MCPServer{}
called := false
middleware := func(next TypedToolHandler[echoArgs]) TypedToolHandler[echoArgs] {
return func(ctx context.Context, headers Headers, args echoArgs) (mcptypes.CallToolResult, error) {
if args.Message != "hello" {
t.Fatalf("middleware args = %#v", args)
}
called = true
return next(ctx, headers, args)
}
}
RegisterTypedToolWithMiddleware(srv, mcptypes.Tool{Name: "echo"}, func(context.Context, Headers, echoArgs) (mcptypes.CallToolResult, error) {
return mcptypes.NewToolResultText("ok"), nil
}, middleware)
req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"echo","arguments":{"message":"hello"}}}`))
req.Header.Set("Mcp-Method", "tools/call")
req.Header.Set("Mcp-Name", "echo")
srv.ServeHTTP(httptest.NewRecorder(), req)
if !called {
t.Fatal("typed middleware was not called")
}
}
func TestServeHTTPContainsToolPanic(t *testing.T) {
srv := &MCPServer{}
srv.RegisterTool(mcptypes.Tool{Name: "panic"}, func(context.Context, CallToolRequest) (mcptypes.CallToolResult, error) {
panic("boom")
})
req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"panic"}}`))
req.Header.Set("Mcp-Method", "tools/call")
req.Header.Set("Mcp-Name", "panic")
res := httptest.NewRecorder()
srv.ServeHTTP(res, req)
var response struct {
Result mcptypes.CallToolResult `json:"result"`
}
if err := json.NewDecoder(res.Body).Decode(&response); err != nil {
t.Fatal(err)
}
if !response.Result.IsError {
t.Fatal("panic did not become an in-band tool error")
}
}
func TestServeHTTPRejectsInvalidTypedArguments(t *testing.T) {
srv := &MCPServer{}
RegisterTypedTool(srv, mcptypes.Tool{Name: "echo"}, func(context.Context, Headers, echoArgs) (mcptypes.CallToolResult, error) {
@ -68,6 +113,27 @@ func TestServeHTTPRejectsInvalidTypedArguments(t *testing.T) {
}
}
func TestRegisterResourceRejectsInvalidURI(t *testing.T) {
defer func() {
if recover() == nil {
t.Fatal("invalid resource URI did not panic")
}
}()
(&MCPServer{}).RegisterResource(mcptypes.Resource{URI: "not absolute"}, nil)
}
func TestServeHTTPRejectsAppOnlyToolsFromList(t *testing.T) {
srv := &MCPServer{}
srv.RegisterTool(mcptypes.Tool{Name: "hidden", Meta: &mcptypes.Meta{UI: &mcptypes.AppUI{Visibility: []mcptypes.Visibility{mcptypes.VisibilityApp}}}}, nil)
req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))
req.Header.Set("Mcp-Method", "tools/list")
res := httptest.NewRecorder()
srv.ServeHTTP(res, req)
if strings.Contains(res.Body.String(), "hidden") {
t.Fatalf("app-only tool leaked into list: %s", res.Body)
}
}
func TestServeHTTPRejectsHeaderNameMismatch(t *testing.T) {
srv := &MCPServer{}
srv.RegisterTool(mcptypes.Tool{Name: "echo"}, func(context.Context, CallToolRequest) (mcptypes.CallToolResult, error) {

149
mcp/types20260728/types.go
View File

@ -10,6 +10,7 @@ import (
"encoding/json"
"errors"
"fmt"
"maps"
)
const (
@ -88,9 +89,125 @@ type RequestMeta struct {
}
type ResultMeta struct {
ServerInfo *ServerInfo `json:"io.modelcontextprotocol/serverInfo,omitempty"`
SubscriptionID string `json:"io.modelcontextprotocol/subscriptionId,omitempty"`
Extensions map[string]json.RawMessage `json:"-"`
ServerInfo *ServerInfo
SubscriptionID string
UI *AppUI
Extensions map[string]json.RawMessage
}
func (m ResultMeta) MarshalJSON() ([]byte, error) {
values := make(map[string]json.RawMessage, len(m.Extensions)+3)
maps.Copy(values, m.Extensions)
if m.ServerInfo != nil {
value, err := json.Marshal(m.ServerInfo)
if err != nil {
return nil, err
}
values["io.modelcontextprotocol/serverInfo"] = value
}
if m.SubscriptionID != "" {
value, err := json.Marshal(m.SubscriptionID)
if err != nil {
return nil, err
}
values["io.modelcontextprotocol/subscriptionId"] = value
}
if m.UI != nil {
value, err := json.Marshal(m.UI)
if err != nil {
return nil, err
}
values["ui"] = value
}
return json.Marshal(values)
}
func (m *ResultMeta) UnmarshalJSON(data []byte) error {
var values map[string]json.RawMessage
if err := json.Unmarshal(data, &values); err != nil {
return err
}
m.Extensions = make(map[string]json.RawMessage, len(values))
for key, value := range values {
switch key {
case "io.modelcontextprotocol/serverInfo":
if err := json.Unmarshal(value, &m.ServerInfo); err != nil {
return err
}
case "io.modelcontextprotocol/subscriptionId":
if err := json.Unmarshal(value, &m.SubscriptionID); err != nil {
return err
}
case "ui":
if err := json.Unmarshal(value, &m.UI); err != nil {
return err
}
default:
m.Extensions[key] = value
}
}
return nil
}
// Meta contains typed MCP Apps metadata plus preserved extension fields.
type Meta struct {
UI *AppUI
Extensions map[string]json.RawMessage
}
type AppUI struct {
ResourceURI string `json:"resourceUri,omitempty"`
Visibility []Visibility `json:"visibility,omitempty"`
CSP *AppCSP `json:"csp,omitempty"`
Permissions *AppPermissions `json:"permissions,omitempty"`
}
type Visibility string
const (
VisibilityModel Visibility = "model"
VisibilityApp Visibility = "app"
)
type AppCSP struct {
ConnectDomains []string `json:"connectDomains,omitempty"`
ResourceDomains []string `json:"resourceDomains,omitempty"`
FrameDomains []string `json:"frameDomains,omitempty"`
BaseURIDomains []string `json:"baseUriDomains,omitempty"`
}
type AppPermission struct{}
type AppPermissions struct {
ClipboardWrite *AppPermission `json:"clipboardWrite,omitempty"`
Camera *AppPermission `json:"camera,omitempty"`
Microphone *AppPermission `json:"microphone,omitempty"`
Geolocation *AppPermission `json:"geolocation,omitempty"`
}
func (m Meta) MarshalJSON() ([]byte, error) {
values := make(map[string]json.RawMessage, len(m.Extensions)+1)
maps.Copy(values, m.Extensions)
if m.UI != nil {
value, err := json.Marshal(m.UI)
if err != nil {
return nil, err
}
values["ui"] = value
}
return json.Marshal(values)
}
func (m *Meta) UnmarshalJSON(data []byte) error {
var values map[string]json.RawMessage
if err := json.Unmarshal(data, &values); err != nil {
return err
}
m.Extensions = make(map[string]json.RawMessage, len(values))
if raw, ok := values["ui"]; ok {
if err := json.Unmarshal(raw, &m.UI); err != nil {
return err
}
delete(values, "ui")
}
maps.Copy(m.Extensions, values)
return nil
}
type Response struct {
@ -209,7 +326,7 @@ const (
)
type Tool struct {
Meta json.RawMessage `json:"_meta,omitempty"`
Meta *Meta `json:"_meta,omitempty"`
Name string `json:"name"`
Title string `json:"title,omitempty"`
Description string `json:"description,omitempty"`
@ -225,20 +342,20 @@ type ToolAnnotations struct {
OpenWorldHint *bool `json:"openWorldHint,omitempty"`
}
type Resource struct {
Meta json.RawMessage `json:"_meta,omitempty"`
URI string `json:"uri"`
Name string `json:"name"`
Title string `json:"title,omitempty"`
Description string `json:"description,omitempty"`
MIMEType string `json:"mimeType,omitempty"`
Size *int64 `json:"size,omitempty"`
Meta *Meta `json:"_meta,omitempty"`
URI string `json:"uri"`
Name string `json:"name"`
Title string `json:"title,omitempty"`
Description string `json:"description,omitempty"`
MIMEType string `json:"mimeType,omitempty"`
Size *int64 `json:"size,omitempty"`
}
type ResourceContents struct {
Meta json.RawMessage `json:"_meta,omitempty"`
URI string `json:"uri"`
MIMEType string `json:"mimeType,omitempty"`
Text string `json:"text,omitempty"`
Blob string `json:"blob,omitempty"`
Meta *Meta `json:"_meta,omitempty"`
URI string `json:"uri"`
MIMEType string `json:"mimeType,omitempty"`
Text string `json:"text,omitempty"`
Blob string `json:"blob,omitempty"`
}
type Content struct {
Type string `json:"type"`

28
mcp/types20260728/types_test.go
View File

@ -10,6 +10,34 @@ type typedResult struct {
Count int `json:"count"`
}
func TestAppsMetadataUsesEmptyPermissionObjects(t *testing.T) {
data, err := json.Marshal(Meta{UI: &AppUI{Permissions: &AppPermissions{ClipboardWrite: &AppPermission{}}}})
if err != nil {
t.Fatal(err)
}
if string(data) != `{"ui":{"permissions":{"clipboardWrite":{}}}}` {
t.Fatalf("metadata = %s", data)
}
}
func TestResultMetadataPreservesAppsUI(t *testing.T) {
original := ResultMeta{UI: &AppUI{ResourceURI: "ui://reports"}, Extensions: map[string]json.RawMessage{"example/trace": json.RawMessage(`"abc"`)}}
data, err := json.Marshal(original)
if err != nil {
t.Fatal(err)
}
var decoded ResultMeta
if err := json.Unmarshal(data, &decoded); err != nil {
t.Fatal(err)
}
if decoded.UI == nil || decoded.UI.ResourceURI != "ui://reports" {
t.Fatalf("UI = %#v", decoded.UI)
}
if string(decoded.Extensions["example/trace"]) != `"abc"` {
t.Fatalf("extensions = %#v", decoded.Extensions)
}
}
func TestNewToolResultStructuredUsesTypedValue(t *testing.T) {
result, err := NewToolResultStructured(typedResult{Step: "done", Count: 2}, "finished")
if err != nil {