go-wikijs-mcp/internal/transport/http.go

// Package transport wires the MCP server to stdio or streamable HTTP.
package transport

import (
	"context"
	"crypto/subtle"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// DefaultSessionTimeout closes idle HTTP sessions after 30 minutes.
const DefaultSessionTimeout = 30 * time.Minute

// HTTPOptions configures the streamable HTTP transport.
type HTTPOptions struct {
	Addr           string
	Path           string
	Stateless      bool
	AuthToken      string
	SessionTimeout time.Duration
	Logger         *slog.Logger
	ProtocolLogger *slog.Logger
	NewServer      func(*http.Request) *mcp.Server
}

// NewHTTPHandler returns an HTTP handler that serves the MCP endpoint at
// opts.Path and a liveness endpoint at /healthz. If AuthToken is set, every
// request except /healthz requires a matching bearer token.
func NewHTTPHandler(opts HTTPOptions) http.Handler {
	if opts.NewServer == nil {
		panic("transport: NewHTTPHandler requires NewServer")
	}
	logger := opts.Logger
	if logger == nil {
		logger = slog.New(slog.DiscardHandler)
	}
	if opts.Path == "" {
		opts.Path = "/mcp"
	}
	if opts.SessionTimeout <= 0 {
		opts.SessionTimeout = DefaultSessionTimeout
	}
	protocolLogger := opts.ProtocolLogger
	if protocolLogger == nil {
		protocolLogger = logger
	}

	mcpHandler := mcp.NewStreamableHTTPHandler(opts.NewServer, &mcp.StreamableHTTPOptions{
		Stateless:      opts.Stateless,
		Logger:         protocolLogger,
		SessionTimeout: opts.SessionTimeout,
	})

	mux := http.NewServeMux()
	mux.Handle(opts.Path, mcpHandler)
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "no-store")
		_, _ = w.Write([]byte(`{"status":"ok"}`))
	})

	var handler http.Handler = mux
	if opts.AuthToken != "" {
		handler = bearerAuth(opts.AuthToken, handler)
	}
	return logRequests(logger, handler)
}

// ServeHTTP runs the streamable HTTP transport until ctx is cancelled.
func ServeHTTP(ctx context.Context, opts HTTPOptions) error {
	logger := opts.Logger
	if logger == nil {
		logger = slog.New(slog.DiscardHandler)
	}

	ln, err := net.Listen("tcp", opts.Addr)
	if err != nil {
		return fmt.Errorf("listen on %s: %w", opts.Addr, err)
	}
	if opts.AuthToken == "" && !isLoopbackListener(ln.Addr()) {
		_ = ln.Close()
		return fmt.Errorf("refusing to serve HTTP on non-loopback address %s without MCP_AUTH_TOKEN", ln.Addr())
	}

	srv := &http.Server{
		Handler:           NewHTTPHandler(opts),
		ReadHeaderTimeout: 10 * time.Second,
		IdleTimeout:       2 * time.Minute,
	}

	logger.Info("WikiJS MCP server listening",
		"transport", "http",
		"addr", ln.Addr().String(),
		"path", opts.Path,
		"stateless", opts.Stateless,
		"auth", opts.AuthToken != "",
	)

	errCh := make(chan error, 1)
	go func() {
		errCh <- srv.Serve(ln)
	}()

	select {
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if err := srv.Shutdown(shutdownCtx); err != nil {
			logger.Warn("HTTP shutdown incomplete, forcing close", "error", err)
			_ = srv.Close()
		}
		return nil
	case err := <-errCh:
		if errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return err
	}
}

func isLoopbackListener(addr net.Addr) bool {
	tcp, ok := addr.(*net.TCPAddr)
	return ok && tcp.IP.IsLoopback()
}

func bearerAuth(token string, next http.Handler) http.Handler {
	expected := []byte(token)
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/healthz" {
			next.ServeHTTP(w, r)
			return
		}
		scheme, got, ok := strings.Cut(r.Header.Get("Authorization"), " ")
		if !ok || !strings.EqualFold(scheme, "Bearer") || subtle.ConstantTimeCompare([]byte(got), expected) != 1 {
			w.Header().Set("WWW-Authenticate", `Bearer realm="wikijs-mcp"`)
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

type statusWriter struct {
	http.ResponseWriter
	status int
}

func (w *statusWriter) WriteHeader(code int) {
	if w.status == 0 {
		w.status = code
	}
	w.ResponseWriter.WriteHeader(code)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	return w.ResponseWriter.Write(b)
}

// Unwrap lets http.ResponseController reach the underlying writer, which is
// required for HTTP streaming.
func (w *statusWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }

func logRequests(logger *slog.Logger, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		sw := &statusWriter{ResponseWriter: w}
		next.ServeHTTP(sw, r)
		status := sw.status
		if status == 0 {
			status = http.StatusOK
		}
		logger.Debug("http request",
			"method", r.Method,
			"path", r.URL.Path,
			"status", status,
			"duration", time.Since(start),
			"remote", r.RemoteAddr,
		)
	})
}