261 lines
8.0 KiB
Go
261 lines
8.0 KiB
Go
package mcpclient
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
"github.com/zalando/go-keyring"
|
|
)
|
|
|
|
type memorySecrets map[string]string
|
|
|
|
func (m memorySecrets) Get(id string) (string, error) {
|
|
if v, ok := m[id]; ok {
|
|
return v, nil
|
|
}
|
|
return "", keyring.ErrNotFound
|
|
}
|
|
func (m memorySecrets) Set(id, value string) error { m[id] = value; return nil }
|
|
func (m memorySecrets) Delete(id string) error { delete(m, id); return nil }
|
|
func testService(t *testing.T) (*Service, memorySecrets) {
|
|
t.Helper()
|
|
store := memorySecrets{}
|
|
return newService(filepath.Join(t.TempDir(), "mcp.json"), store), store
|
|
}
|
|
func httpInput() ServerInput {
|
|
return ServerInput{Config: ServerConfig{Name: "Test", Enabled: true, Transport: "http", URL: "http://127.0.0.1:8080/mcp", TimeoutSeconds: 3}}
|
|
}
|
|
|
|
func TestPersistenceAndSecretLifecycle(t *testing.T) {
|
|
s, secrets := testService(t)
|
|
in := httpInput()
|
|
in.ReplaceSecrets = true
|
|
in.Secrets.Headers = map[string]string{"Authorization": "Bearer PRIVATE-TOKEN"}
|
|
view, err := s.SaveServer(in)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !view.HasSecrets || view.Config.ID == "" {
|
|
t.Fatalf("invalid view: %+v", view)
|
|
}
|
|
file, err := os.ReadFile(s.path)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
views, err := s.ListServers()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
public, _ := json.Marshal(views)
|
|
if strings.Contains(string(file)+string(public), "PRIVATE-TOKEN") {
|
|
t.Fatal("secret exposed")
|
|
}
|
|
if err := s.SetEnabled(view.Config.ID, false); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
reloaded := newService(s.path, secrets)
|
|
views, err = reloaded.ListServers()
|
|
if err != nil || len(views) != 1 || views[0].Config.Enabled || !views[0].HasSecrets {
|
|
t.Fatalf("reload: %v %+v", err, views)
|
|
}
|
|
in = ServerInput{Config: views[0].Config}
|
|
in.Config.Name = "Renamed"
|
|
if _, err := reloaded.SaveServer(in); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(secrets) != 1 {
|
|
t.Fatal("existing secret lost")
|
|
}
|
|
in.Config.URL = "http://127.0.0.1:9090/mcp"
|
|
if _, err := reloaded.SaveServer(in); err == nil {
|
|
t.Fatal("retarget reused credentials")
|
|
}
|
|
in.ReplaceSecrets = true
|
|
if v, err := reloaded.SaveServer(in); err != nil || v.HasSecrets {
|
|
t.Fatalf("clear: %v %+v", err, v)
|
|
}
|
|
if len(secrets) != 0 {
|
|
t.Fatal("secret not cleared")
|
|
}
|
|
if err := reloaded.DeleteServer(view.Config.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
views, err = newService(s.path, secrets).ListServers()
|
|
if err != nil || len(views) != 0 {
|
|
t.Fatalf("delete persisted: %v %+v", err, views)
|
|
}
|
|
}
|
|
|
|
func TestFailedWriteRollsBackCredentials(t *testing.T) {
|
|
s, secrets := testService(t)
|
|
in := httpInput()
|
|
in.ReplaceSecrets = true
|
|
in.Secrets.Headers = map[string]string{"Authorization": "original"}
|
|
view, err := s.SaveServer(in)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
original := secrets[view.Config.ID]
|
|
// Renaming a temporary file onto a directory must fail on all platforms.
|
|
s.path = t.TempDir()
|
|
in.Config = view.Config
|
|
in.Secrets.Headers["Authorization"] = "replacement"
|
|
if _, err := s.SaveServer(in); err == nil {
|
|
t.Fatal("write unexpectedly succeeded")
|
|
}
|
|
if secrets[view.Config.ID] != original {
|
|
t.Fatal("save did not roll back secret")
|
|
}
|
|
if err := s.DeleteServer(view.Config.ID); err == nil {
|
|
t.Fatal("delete unexpectedly succeeded")
|
|
}
|
|
if secrets[view.Config.ID] != original || len(s.servers) != 1 {
|
|
t.Fatal("delete did not roll back")
|
|
}
|
|
}
|
|
|
|
func TestValidationAndCorruptConfiguration(t *testing.T) {
|
|
s, _ := testService(t)
|
|
cases := []func(*ServerInput){
|
|
func(in *ServerInput) { in.Config.URL = "http://public.example/mcp" },
|
|
func(in *ServerInput) { in.Config.URL = "https://user:password@example.com/mcp" },
|
|
func(in *ServerInput) { in.Config.TimeoutSeconds = 0 },
|
|
func(in *ServerInput) { in.Config.Transport = "sse" },
|
|
func(in *ServerInput) { in.Config.ID = "unknown" },
|
|
func(in *ServerInput) {
|
|
in.ReplaceSecrets = true
|
|
in.Secrets.Headers = map[string]string{"Authorization": "a\r\nb"}
|
|
},
|
|
func(in *ServerInput) {
|
|
in.ReplaceSecrets = true
|
|
in.Secrets.Headers = map[string]string{"Mcp-Session-Id": "override"}
|
|
},
|
|
}
|
|
for i, change := range cases {
|
|
in := httpInput()
|
|
change(&in)
|
|
if _, err := s.SaveServer(in); err == nil {
|
|
t.Errorf("case %d accepted", i)
|
|
}
|
|
}
|
|
if err := os.WriteFile(s.path, []byte("{broken"), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
broken := newService(s.path, memorySecrets{})
|
|
if _, err := broken.ListServers(); err == nil {
|
|
t.Fatal("corruption hidden")
|
|
}
|
|
if _, err := broken.SaveServer(httpInput()); err == nil {
|
|
t.Fatal("corrupt config overwritten")
|
|
}
|
|
}
|
|
|
|
func newTestMCPServer(calls *atomic.Int32) *mcp.Server {
|
|
server := mcp.NewServer(&mcp.Implementation{Name: "fixture", Version: "1"}, &mcp.ServerOptions{PageSize: 1})
|
|
for _, name := range []string{"alpha", "beta"} {
|
|
mcp.AddTool(server, &mcp.Tool{Name: name, Description: "A test tool"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) {
|
|
calls.Add(1)
|
|
return &mcp.CallToolResult{}, nil, nil
|
|
})
|
|
}
|
|
return server
|
|
}
|
|
|
|
func TestHTTPDiscoveryAuthPaginationAndNoToolExecution(t *testing.T) {
|
|
var calls, authenticated atomic.Int32
|
|
server := newTestMCPServer(&calls)
|
|
handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return server }, nil)
|
|
remote := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Header.Get("Authorization") != "Bearer fixture-token" {
|
|
http.Error(w, "unauthorized fixture-token", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
authenticated.Add(1)
|
|
handler.ServeHTTP(w, r)
|
|
}))
|
|
defer remote.Close()
|
|
s, _ := testService(t)
|
|
in := httpInput()
|
|
in.Config.URL = remote.URL
|
|
in.ReplaceSecrets = true
|
|
in.Secrets.Headers = map[string]string{"Authorization": "Bearer fixture-token"}
|
|
result := s.TestServer(context.Background(), in)
|
|
if !result.OK || len(result.Tools) != 2 || !strings.Contains(result.Tools[0].InputSchema, "object") {
|
|
t.Fatalf("discovery: %+v", result)
|
|
}
|
|
if calls.Load() != 0 || authenticated.Load() < 2 {
|
|
t.Fatalf("calls=%d authenticated=%d", calls.Load(), authenticated.Load())
|
|
}
|
|
if _, err := os.Stat(s.path); !errors.Is(err, os.ErrNotExist) {
|
|
t.Fatal("test persisted draft")
|
|
}
|
|
in.Secrets.Headers = nil
|
|
result = s.TestServer(context.Background(), in)
|
|
if result.OK || strings.Contains(result.Message, "fixture-token") || len(result.Tools) != 0 {
|
|
t.Fatalf("auth failure: %+v", result)
|
|
}
|
|
}
|
|
|
|
func TestDiscoveryCancellation(t *testing.T) {
|
|
release := make(chan struct{})
|
|
remote := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
select {
|
|
case <-r.Context().Done():
|
|
case <-release:
|
|
}
|
|
}))
|
|
defer remote.Close()
|
|
defer close(release)
|
|
s, _ := testService(t)
|
|
in := httpInput()
|
|
in.Config.URL = remote.URL
|
|
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
|
defer cancel()
|
|
start := time.Now()
|
|
result := s.TestServer(ctx, in)
|
|
if result.OK || time.Since(start) > 2*time.Second {
|
|
t.Fatalf("cancellation: %+v", result)
|
|
}
|
|
}
|
|
|
|
func TestStdioDiscovery(t *testing.T) {
|
|
s, _ := testService(t)
|
|
executable, err := os.Executable()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
in := ServerInput{Config: ServerConfig{Name: "Stdio", Transport: "stdio", Command: executable, Args: []string{"-test.run=^TestMCPHelperProcess$"}, WorkingDirectory: t.TempDir(), TimeoutSeconds: 5}, ReplaceSecrets: true, Secrets: Secrets{Env: map[string]string{"ITUOKE_MCP_HELPER": "1", "MCP_TEST_SECRET": "helper-token"}}}
|
|
result := s.TestServer(context.Background(), in)
|
|
if !result.OK || len(result.Tools) != 2 {
|
|
t.Fatalf("stdio: %+v", result)
|
|
}
|
|
in.Config.Command = filepath.Join(t.TempDir(), "missing-executable")
|
|
if result = s.TestServer(context.Background(), in); result.OK {
|
|
t.Fatal("missing executable accepted")
|
|
}
|
|
}
|
|
|
|
func TestMCPHelperProcess(t *testing.T) {
|
|
if os.Getenv("ITUOKE_MCP_HELPER") != "1" {
|
|
return
|
|
}
|
|
if os.Getenv("MCP_TEST_SECRET") != "helper-token" {
|
|
os.Exit(2)
|
|
}
|
|
var calls atomic.Int32
|
|
if err := newTestMCPServer(&calls).Run(context.Background(), &mcp.StdioTransport{}); err != nil {
|
|
os.Exit(3)
|
|
}
|
|
os.Exit(0)
|
|
}
|