From 52ee285b182cbea1bf39e24e88fe676f9d52f65e Mon Sep 17 00:00:00 2001
From: cryo <zdtna412@gmail.com>
Date: Wed, 6 Aug 2025 16:09:56 +0000
Subject: [PATCH] mcp: add syntax and scheme validation to AddResourceTemplate

---
 mcp/server.go      | 15 ++++++++++++++-
 mcp/server_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/mcp/server.go b/mcp/server.go
index c8878da3..75cde63c 100644
--- a/mcp/server.go
+++ b/mcp/server.go
@@ -22,6 +22,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/internal/util"
 	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
+	"github.com/yosida95/uritemplate/v3"
 )
 
 const DefaultPageSize = 1000
@@ -220,7 +221,19 @@ func (s *Server) RemoveResources(uris ...string) {
 func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) {
 	s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
 		func() bool {
-			// TODO: check template validity.
+			// Validate the URI template syntax
+			_, err := uritemplate.New(t.URITemplate)
+			if err != nil {
+				panic(fmt.Errorf("URI template %s is invalid: %v", t.URITemplate, err))
+			}
+			// Ensure the URI template has a valid scheme
+			u, err := url.Parse(t.URITemplate)
+			if err != nil {
+				panic(err) // url.Parse includes the URI in the error
+			}
+			if !u.IsAbs() {
+				panic(fmt.Errorf("URI template %s needs a scheme", t.URITemplate))
+			}
 			s.resourceTemplates.add(&serverResourceTemplate{t, h})
 			return true
 		})
diff --git a/mcp/server_test.go b/mcp/server_test.go
index 5a161b72..8d904da8 100644
--- a/mcp/server_test.go
+++ b/mcp/server_test.go
@@ -371,3 +371,39 @@ func TestServerCapabilities(t *testing.T) {
 		})
 	}
 }
+
+func TestServerAddResourceTemplate(t *testing.T) {
+	tests := []struct {
+		name        string
+		template    string
+		expectPanic bool
+	}{
+		{"ValidFileTemplate", "file:///{a}/{b}", false},
+		{"ValidCustomScheme", "myproto:///{a}", false},
+		{"MissingScheme1", "://example.com/{path}", true},
+		{"MissingScheme2", "/api/v1/users/{id}", true},
+		{"EmptyVariable", "file:///{}/{b}", true},
+		{"UnclosedVariable", "file:///{a", true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rt := ResourceTemplate{URITemplate: tt.template}
+
+			defer func() {
+				if r := recover(); r != nil {
+					if !tt.expectPanic {
+						t.Errorf("%s: unexpected panic: %v", tt.name, r)
+					}
+				} else {
+					if tt.expectPanic {
+						t.Errorf("%s: expected panic but did not panic", tt.name)
+					}
+				}
+			}()
+
+			s := NewServer(testImpl, nil)
+			s.AddResourceTemplate(&rt, nil)
+		})
+	}
+}