feat: restore MCP server safety features

mcp/server20260728/server.go

@@ -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))
+	tools := make([]mcptypes.Tool, 0, len(s.tools))
-	for i := range s.tools {
+	for _, registered := range s.tools {
-		tools[i] = s.tools[i].definition
+		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}
 }

mcp/server20260728/server_test.go

@@ -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) {

mcp/types20260728/types.go

@@ -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"`
+	ServerInfo     *ServerInfo
-	SubscriptionID string                     `json:"io.modelcontextprotocol/subscriptionId,omitempty"`
+	SubscriptionID string
-	Extensions     map[string]json.RawMessage `json:"-"`
+	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,7 +342,7 @@ type ToolAnnotations struct {
 	OpenWorldHint   *bool  `json:"openWorldHint,omitempty"`
 }
 type Resource struct {
-	Meta        json.RawMessage `json:"_meta,omitempty"`
+	Meta        *Meta  `json:"_meta,omitempty"`
 	URI         string `json:"uri"`
 	Name        string `json:"name"`
 	Title       string `json:"title,omitempty"`
@@ -234,7 +351,7 @@ type Resource struct {
 	Size        *int64 `json:"size,omitempty"`
 }
 type ResourceContents struct {
-	Meta     json.RawMessage `json:"_meta,omitempty"`
+	Meta     *Meta  `json:"_meta,omitempty"`
 	URI      string `json:"uri"`
 	MIMEType string `json:"mimeType,omitempty"`
 	Text     string `json:"text,omitempty"`

mcp/types20260728/types_test.go

@@ -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 {