sm-bm-plugin/mcp_links.go

package bm

import (
	"bytes"
	"encoding/json"
	"mime"
	"net/http"
	"strconv"
	"strings"

	"git.golem15.com/golem15/summercms/modules/bouncer"
)

// MiddlewareMCPLinks makes the links of an MCP answer absolute (D-16). It
// runs after user.api_token on the token twins of the student API and on the
// management API.
const MiddlewareMCPLinks = "jz.bm.mcp-links"

// mcpLinks rewrites the 2xx JSON answer of an MCP tool call (a request pier
// dispatched in-process with the api_token marker): every string under a key
// named url or ending in _url that starts with a single / gets the
// application's base URL in front, so an agent can follow it. Any other
// request passes straight through, unbuffered and byte for byte.
func (p *Plugin) mcpLinks(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, ok := bouncer.Dispatched(r.Context(), "api_token"); !ok {
			next.ServeHTTP(w, r)
			return
		}
		rec := &linkRecorder{header: http.Header{}}
		next.ServeHTTP(rec, r)
		status := rec.status
		if status == 0 {
			status = http.StatusOK
		}
		out := rec.body.Bytes()
		if status >= 200 && status < 300 && isJSON(rec.header.Get("Content-Type")) && len(out) > 0 {
			if rewritten, ok := absoluteLinks(out, p.linkBase(r)); ok {
				out = rewritten
				rec.header.Set("Content-Length", strconv.Itoa(len(out)))
			}
		}
		for k, v := range rec.header {
			w.Header()[k] = v
		}
		w.WriteHeader(status)
		_, _ = w.Write(out)
	})
}

// linkBase is the base of the absolute links: app.url, else
// jz.bm.public_url, else the scheme (TLS or X-Forwarded-Proto) and Host of
// the request (pier copies both onto the request it dispatches).
func (p *Plugin) linkBase(r *http.Request) string {
	if p.app != nil && p.app.Config != nil {
		for _, key := range []string{"app.url", "jz.bm.public_url"} {
			if base := strings.TrimRight(strings.TrimSpace(p.app.Config.String(key)), "/"); base != "" {
				return base
			}
		}
	}
	scheme := "http"
	if r.TLS != nil || strings.EqualFold(r.Header.Get("X-Forwarded-Proto"), "https") {
		scheme = "https"
	}
	return scheme + "://" + r.Host
}

// isJSON reports whether a Content-Type is JSON.
func isJSON(contentType string) bool {
	mt, _, err := mime.ParseMediaType(contentType)
	return err == nil && (mt == "application/json" || strings.HasSuffix(mt, "+json"))
}

// absoluteLinks decodes body (numbers kept exactly), prefixes its relative
// links with base and encodes it again. It reports false when body is not a
// single JSON value, so the caller sends it unchanged.
func absoluteLinks(body []byte, base string) ([]byte, bool) {
	dec := json.NewDecoder(bytes.NewReader(body))
	dec.UseNumber()
	var v any
	if err := dec.Decode(&v); err != nil {
		return nil, false
	}
	if dec.More() {
		return nil, false
	}
	v = rewriteLinks(v, "", base)
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	enc.SetEscapeHTML(false)
	if err := enc.Encode(v); err != nil {
		return nil, false
	}
	return buf.Bytes(), true
}

// rewriteLinks walks v; key is the object key v sits under ("" inside an
// array or at the top).
func rewriteLinks(v any, key, base string) any {
	switch t := v.(type) {
	case map[string]any:
		for k, child := range t {
			t[k] = rewriteLinks(child, k, base)
		}
		return t
	case []any:
		for i, child := range t {
			t[i] = rewriteLinks(child, "", base)
		}
		return t
	case string:
		if (key == "url" || strings.HasSuffix(key, "_url")) && strings.HasPrefix(t, "/") && !strings.HasPrefix(t, "//") {
			return base + t
		}
		return t
	default:
		return v
	}
}

// linkRecorder buffers a handler's answer so mcpLinks can rewrite it.
type linkRecorder struct {
	header http.Header
	status int
	body   bytes.Buffer
}

func (l *linkRecorder) Header() http.Header { return l.header }

func (l *linkRecorder) WriteHeader(status int) {
	if l.status == 0 {
		l.status = status
	}
}

func (l *linkRecorder) Write(b []byte) (int, error) {
	if l.status == 0 {
		l.status = http.StatusOK
	}
	return l.body.Write(b)
}