From d96fd2087557311f2762a5495889e0e1805a259b Mon Sep 17 00:00:00 2001
From: Sam Morrow <info@sam-morrow.com>
Date: Tue, 18 Mar 2025 15:51:36 +0100
Subject: [PATCH] fix: ensure resource/template/list is marshalled safely

---
 mcp/resources.go      |  2 +-
 mcp/types.go          | 23 ++++++++++++++++++++++-
 server/server.go      |  3 +--
 server/server_test.go | 35 ++++++++++++++++++++++++++++++++---
 4 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/mcp/resources.go b/mcp/resources.go
index ad7d3d560..51cdd25dd 100644
--- a/mcp/resources.go
+++ b/mcp/resources.go
@@ -62,7 +62,7 @@ type ResourceTemplateOption func(*ResourceTemplate)
 // Options are applied in order, allowing for flexible template configuration.
 func NewResourceTemplate(uriTemplate string, name string, opts ...ResourceTemplateOption) ResourceTemplate {
 	template := ResourceTemplate{
-		URITemplate: uritemplate.MustNew(uriTemplate),
+		URITemplate: &URITemplate{Template: uritemplate.MustNew(uriTemplate)},
 		Name:        name,
 	}
 
diff --git a/mcp/types.go b/mcp/types.go
index 2014b531f..c94d58e9c 100644
--- a/mcp/types.go
+++ b/mcp/types.go
@@ -8,6 +8,27 @@ import (
 	"github.com/yosida95/uritemplate/v3"
 )
 
+type URITemplate struct {
+	*uritemplate.Template
+}
+
+func (t *URITemplate) MarshalJSON() ([]byte, error) {
+	return json.Marshal(t.Template.Raw())
+}
+
+func (t *URITemplate) UnmarshalJSON(data []byte) error {
+	var raw string
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+	template, err := uritemplate.New(raw)
+	if err != nil {
+		return err
+	}
+	t.Template = template
+	return nil
+}
+
 /* JSON-RPC types */
 
 // JSONRPCMessage represents either a JSONRPCRequest, JSONRPCNotification, JSONRPCResponse, or JSONRPCError
@@ -446,7 +467,7 @@ type ResourceTemplate struct {
 	Annotated
 	// A URI template (according to RFC 6570) that can be used to construct
 	// resource URIs.
-	URITemplate *uritemplate.Template `json:"uriTemplate"`
+	URITemplate *URITemplate `json:"uriTemplate"`
 	// A human-readable name for the type of resource this template refers to.
 	//
 	// This can be used by clients to populate UI elements.
diff --git a/server/server.go b/server/server.go
index 9bee13b69..5e851a059 100644
--- a/server/server.go
+++ b/server/server.go
@@ -10,7 +10,6 @@ import (
 	"sync/atomic"
 
 	"github.com/mark3labs/mcp-go/mcp"
-	"github.com/yosida95/uritemplate/v3"
 )
 
 // resourceEntry holds both a resource and its handler
@@ -731,7 +730,7 @@ func (s *MCPServer) handleReadResource(
 }
 
 // matchesTemplate checks if a URI matches a URI template pattern
-func matchesTemplate(uri string, template *uritemplate.Template) bool {
+func matchesTemplate(uri string, template *mcp.URITemplate) bool {
 	return template.Regexp().MatchString(uri)
 }
 
diff --git a/server/server_test.go b/server/server_test.go
index 3ff5ce7e7..c0c2ad30f 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -923,7 +923,6 @@ func TestMCPServer_Instructions(t *testing.T) {
 func TestMCPServer_ResourceTemplates(t *testing.T) {
 	server := NewMCPServer("test-server", "1.0.0",
 		WithResourceCapabilities(true, true),
-		WithPromptCapabilities(true),
 	)
 
 	server.AddResourceTemplate(
@@ -947,9 +946,15 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
 		},
 	)
 
-	message := `{
+	listMessage := `{
 		"jsonrpc": "2.0",
 		"id": 1,
+		"method": "resources/templates/list"
+	}`
+
+	message := `{
+		"jsonrpc": "2.0",
+		"id": 2,
 		"method": "resources/read",
 		"params": {
 			"uri": "test://something/test-resource/a/b/c"
@@ -959,12 +964,35 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
 	t.Run("Get resource template", func(t *testing.T) {
 		response := server.HandleMessage(
 			context.Background(),
-			[]byte(message),
+			[]byte(listMessage),
 		)
 		assert.NotNil(t, response)
 
 		resp, ok := response.(mcp.JSONRPCResponse)
 		assert.True(t, ok)
+		listResult, ok := resp.Result.(mcp.ListResourceTemplatesResult)
+		assert.True(t, ok)
+		assert.Len(t, listResult.ResourceTemplates, 1)
+		assert.Equal(t, "My Resource", listResult.ResourceTemplates[0].Name)
+		template, err := json.Marshal(listResult.ResourceTemplates[0])
+		assert.NoError(t, err)
+
+		// Need to serialize the json to map[string]string to validate the URITemplate is correctly marshalled
+		var resourceTemplate map[string]string
+		err = json.Unmarshal(template, &resourceTemplate)
+		assert.NoError(t, err)
+
+		assert.Equal(t, "test://{a}/test-resource{/b*}", resourceTemplate["uriTemplate"])
+
+		response = server.HandleMessage(
+			context.Background(),
+			[]byte(message),
+		)
+
+		assert.NotNil(t, response)
+
+		resp, ok = response.(mcp.JSONRPCResponse)
+		assert.True(t, ok)
 		// Validate that the resource values are returned correctly
 		result, ok := resp.Result.(mcp.ReadResourceResult)
 		assert.True(t, ok)
@@ -974,6 +1002,7 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
 		assert.Equal(t, "test://something/test-resource/a/b/c", resultContent.URI)
 		assert.Equal(t, "text/plain", resultContent.MIMEType)
 		assert.Equal(t, "test content: something", resultContent.Text)
+
 	})
 }