verae-mcp-timestamping-server-go/internal/server/public_test.go

package server_test

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"

	"git.georgelambert.org/marchon/verae-mcp-timestamping-server-go/internal/ledger"
	"git.georgelambert.org/marchon/verae-mcp-timestamping-server-go/internal/server"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// fakeVT records the bearer of every upstream call. "good-key" is a valid credential; the host login is "host-jwt".
type fakeVT struct {
	srv     *httptest.Server
	stamps  atomic.Int32
	hostUse atomic.Int32
	users   atomic.Int32
	batches atomic.Int32
}

func newFakeVT(t *testing.T) *fakeVT {
	f := &fakeVT{}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		if auth == "Bearer host-jwt" {
			f.hostUse.Add(1)
		}
		ok := auth == "Bearer good-key" || auth == "Bearer host-jwt"
		w.Header().Set("Content-Type", "application/json")
		switch {
		case strings.HasPrefix(r.URL.Path, "/auth/login"):
			json.NewEncoder(w).Encode(map[string]any{"token": "host-jwt", "username": "svc", "role": "admin", "expires_at": "2099-01-01T00:00:00Z"})
		case r.URL.Path == "/auth/signup":
			// Self-serve signup is unauthenticated on VeraeTime.
			f.users.Add(1)
			json.NewEncoder(w).Encode(map[string]any{"id": "u1", "username": "x", "role": "user"})
		case !ok:
			w.WriteHeader(http.StatusUnauthorized)
			json.NewEncoder(w).Encode(map[string]any{"error": "unauthorized"})
		case r.URL.Path == "/api/status/mcp-credential-check":
			w.WriteHeader(http.StatusNotFound)
			json.NewEncoder(w).Encode(map[string]any{"error": "job not found"})
		case r.URL.Path == "/api/batch/timestamp":
			var body struct {
				Items []map[string]any `json:"items"`
			}
			_ = json.NewDecoder(r.Body).Decode(&body)
			if len(body.Items) == 1 && body.Items[0]["digest"] == strings.Repeat("c", 64) {
				w.WriteHeader(http.StatusPaymentRequired)
				json.NewEncoder(w).Encode(map[string]any{"error": "payment required: Insufficient prepaid balance: this request costs 10 credits but only 0 are available", "code": "PAYMENT_REQUIRED"})
				return
			}
			f.batches.Add(1)
			res := []map[string]any{}
			for i, it := range body.Items {
				res = append(res, map[string]any{"index": i, "digest": it["digest"], "jobId": "job-" + string(rune('a'+i)), "status": "submitted", "existing": false, "credits": 10})
			}
			json.NewEncoder(w).Encode(map[string]any{"count": len(res), "created": len(res), "existing": 0, "failed": 0, "totalCredits": 10 * len(res), "balanceCredits": 480, "results": res})
		case r.URL.Path == "/api/timestamp":
			f.stamps.Add(1)
			json.NewEncoder(w).Encode(map[string]any{"job_id": "j1", "status": "completed", "certificate": "QmTestCid", "merkle_root": "ab", "created_at": "2026-10-10T00:00:00Z"})
		case strings.HasPrefix(r.URL.Path, "/api/status/"):
			json.NewEncoder(w).Encode(map[string]any{"job_id": "j1", "status": "completed", "certificate": "QmTestCid", "merkle_root": "ab"})
		case r.URL.Path == "/api/verify":
			json.NewEncoder(w).Encode(map[string]any{"valid": true, "digest": strings.Repeat("a", 64)})
		case strings.HasPrefix(r.URL.Path, "/api/hash/"):
			json.NewEncoder(w).Encode(map[string]any{"receipts": []any{}})
		case r.URL.Path == "/auth/validate":
			json.NewEncoder(w).Encode(map[string]any{"valid": true, "username": "who", "role": "user"})
		case r.URL.Path == "/auth/users":
			f.users.Add(1)
			json.NewEncoder(w).Encode(map[string]any{"id": "u1", "username": "x", "role": "user"})
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	t.Cleanup(f.srv.Close)
	return f
}

func publicSession(t *testing.T, vt *fakeVT, bearer, xff string) *mcp.ClientSession {
	t.Helper()
	dir := t.TempDir()
	led, err := ledger.Open(filepath.Join(dir, "ledger.json"))
	if err != nil {
		t.Fatal(err)
	}
	authPath := filepath.Join(dir, "auth.properties")
	if err := os.WriteFile(authPath, []byte("username=svc\npassword=pw\ntoken=host-jwt\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	t.Setenv("VERAE_GLOBAL_TOKEN", "host-jwt")
	srv, err := server.New(server.Options{Ledger: led, AuthPath: authPath, GlobalURL: vt.srv.URL, JSONResponse: true, Public: true})
	if err != nil {
		t.Fatal(err)
	}
	h := srv.Handler()
	hs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if bearer != "" {
			r.Header.Set("Authorization", "Bearer "+bearer)
		}
		if xff != "" {
			r.Header.Set("X-Forwarded-For", xff)
		}
		h.ServeHTTP(w, r)
	}))
	t.Cleanup(hs.Close)
	c := mcp.NewClient(&mcp.Implementation{Name: "t", Version: "0"}, nil)
	s, err := c.Connect(context.Background(), &mcp.StreamableClientTransport{Endpoint: hs.URL + "/mcp"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s
}

func call(t *testing.T, s *mcp.ClientSession, name string, args map[string]any) (bool, string) {
	t.Helper()
	res, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: name, Arguments: args})
	if err != nil {
		return false, err.Error()
	}
	b, _ := json.Marshal(res)
	return !res.IsError, string(b)
}

var dig = strings.Repeat("a", 64)

func TestPublicAnonymousWritesRefusedNoHostFallback(t *testing.T) {
	vt := newFakeVT(t)
	s := publicSession(t, vt, "", "203.0.113.5")
	for _, c := range []struct {
		tool string
		args map[string]any
	}{
		{"stamp", map[string]any{"sha256": dig, "scope": "global"}},
		{"stamp", map[string]any{"sha256": dig}}, // local ledger on the server is a write too
		{"create_user", map[string]any{"username": "evil", "password": "pw123456", "role": "admin"}},
		{"login", map[string]any{"username": "admin", "password": "x"}},
	} {
		if ok, out := call(t, s, c.tool, c.args); ok {
			t.Fatalf("%s anonymous succeeded: %s", c.tool, out)
		}
	}
	// Public self-serve signup is allowed without a caller bearer (creates the account only).
	if ok, out := call(t, s, "signup", map[string]any{"username": "spam", "password": "pw123456"}); !ok {
		t.Fatalf("anonymous signup should succeed: %s", out)
	}
	if vt.users.Load() != 1 {
		t.Fatalf("signup should hit /auth/signup once, users=%d", vt.users.Load())
	}
	ok, out := call(t, s, "session", map[string]any{})
	// Anonymous session is a tool error (not a silent valid:false) so agents see how to attach Bearer.
	if ok || strings.Contains(out, `"valid":true`) || !strings.Contains(out, "needs your own Verae credential") {
		t.Fatalf("anonymous session must error with credential guidance: %v %s", ok, out)
	}
	if vt.stamps.Load() != 0 || vt.hostUse.Load() != 0 {
		t.Fatalf("upstream stamps or host-token use: stamps=%d host=%d", vt.stamps.Load(), vt.hostUse.Load())
	}
}

func TestPublicBadKeyRefused(t *testing.T) {
	vt := newFakeVT(t)
	s := publicSession(t, vt, "xVerae-wrong", "203.0.113.6")
	if ok, out := call(t, s, "stamp", map[string]any{"sha256": dig, "scope": "global"}); ok {
		t.Fatalf("bad key stamped: %s", out)
	}
	if ok, _ := call(t, s, "stamp", map[string]any{"sha256": dig}); ok {
		t.Fatal("bad key wrote the local ledger")
	}
	if vt.stamps.Load() != 0 || vt.hostUse.Load() != 0 {
		t.Fatalf("stamps=%d host=%d", vt.stamps.Load(), vt.hostUse.Load())
	}
}

func TestPublicKeyedWritesUseCallerCredential(t *testing.T) {
	vt := newFakeVT(t)
	s := publicSession(t, vt, "good-key", "203.0.113.7")
	if ok, out := call(t, s, "stamp", map[string]any{"sha256": dig, "scope": "global"}); !ok {
		t.Fatalf("keyed global stamp: %s", out)
	}
	if ok, out := call(t, s, "stamp", map[string]any{"sha256": dig}); !ok {
		t.Fatalf("keyed local stamp: %s", out)
	}
	if ok, out := call(t, s, "session", map[string]any{}); !ok || !strings.Contains(out, `"valid":true`) {
		t.Fatalf("keyed session: %s", out)
	}
	if ok, _ := call(t, s, "login", map[string]any{"username": "a", "password": "b"}); ok {
		t.Fatal("login must stay disabled on the public server")
	}
	if vt.stamps.Load() != 1 || vt.hostUse.Load() != 0 {
		t.Fatalf("stamps=%d host=%d", vt.stamps.Load(), vt.hostUse.Load())
	}
}

func TestPublicAnonymousReadsNeverLimited(t *testing.T) {
	vt := newFakeVT(t)
	s := publicSession(t, vt, "", "198.51.100.9")
	// DEC-660: well past the old 30/min (burst 15) per client and in one burst: every anonymous read answers
	for i := 0; i < 200; i++ {
		tool, args := "verify", map[string]any{"receipt_id": "QmTestCid", "want_scope": "global"}
		if i%2 == 1 {
			tool, args = "hash_report", map[string]any{"sha256": dig, "want_scope": "global"}
		}
		if ok, out := call(t, s, tool, args); !ok {
			t.Fatalf("anonymous %s #%d refused: %s", tool, i, out)
		}
	}
	// the write lockdown stays: an anonymous stamp or stamp_batch is still refused
	if ok, out := call(t, s, "stamp", map[string]any{"sha256": dig, "scope": "global"}); ok || !strings.Contains(out, "needs your own Verae credential") {
		t.Fatalf("anonymous stamp: %v %s", ok, out)
	}
	if ok, out := call(t, s, "stamp_batch", map[string]any{"sha256s": []string{dig}, "scope": "global"}); ok || !strings.Contains(out, "needs your own Verae credential") {
		t.Fatalf("anonymous stamp_batch: %v %s", ok, out)
	}
	if vt.stamps.Load() != 0 || vt.batches.Load() != 0 {
		t.Fatal("anonymous write reached VeraeTime")
	}
}

func TestPublicStampBatchWithKey(t *testing.T) {
	vt := newFakeVT(t)
	s := publicSession(t, vt, "good-key", "198.51.100.9")
	ok, out := call(t, s, "stamp_batch", map[string]any{"sha256s": []string{dig, strings.Repeat("b", 64)}, "scope": "global"})
	if !ok || !strings.Contains(out, `"total_credits":20`) || !strings.Contains(out, `"balance_credits":480`) || !strings.Contains(out, `"index":1`) {
		t.Fatalf("stamp_batch: %v %s", ok, out)
	}
	if vt.batches.Load() != 1 || vt.hostUse.Load() != 0 {
		t.Fatalf("batches=%d host=%d", vt.batches.Load(), vt.hostUse.Load())
	}
	// 402 from VeraeTime/zap comes back as a clear message
	ok, out = call(t, s, "stamp_batch", map[string]any{"sha256s": []string{strings.Repeat("c", 64)}, "scope": "global"})
	if ok || !strings.Contains(out, "Insufficient prepaid balance") {
		t.Fatalf("stamp_batch 402: %v %s", ok, out)
	}
}