spark-mcp/internal/mcp/tools/fetch_url.go

package tools

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/mark3labs/mcp-go/mcp"

	"spark-mcp-go/internal/cluster"
	"spark-mcp-go/internal/httpclient"
)

const FetchURLName = "fetch_url"

var fetchURLMethods = map[string]bool{
	http.MethodGet:    true,
	http.MethodPost:   true,
	http.MethodPut:    true,
	http.MethodDelete: true,
	http.MethodHead:   true,
}

// NewFetchURLTool returns the schema for the fetch_url MCP Tool.
func NewFetchURLTool() mcp.Tool {
	return mcp.NewTool(FetchURLName,
		mcp.WithDescription("Perform an HTTP request to an allowed URL using a cluster's auth and allowlist. Returns {status, status_code, headers, body, body_truncated}."),
		mcp.WithString("cluster_id",
			mcp.Required(),
			mcp.Description("ID of the configured cluster whose allowlist and auth are used"),
		),
		mcp.WithString("url",
			mcp.Required(),
			mcp.Description("Absolute http(s) URL to fetch"),
		),
		mcp.WithString("method",
			mcp.Description("HTTP method"),
			mcp.Enum("GET", "POST", "PUT", "DELETE", "HEAD"),
			mcp.DefaultString("GET"),
		),
		mcp.WithObject("headers",
			mcp.Description("Extra HTTP headers as a JSON object"),
			mcp.AdditionalProperties(map[string]any{"type": "string"}),
		),
		mcp.WithString("body",
			mcp.Description("Request body for POST/PUT"),
		),
	)
}

// FetchURLHandler performs an HTTP request with cluster auth and SSRF protection.
func (d *Deps) FetchURLHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	clusterID, err := req.RequireString("cluster_id")
	if err != nil {
		return errResult("fetch_url: " + err.Error()), nil
	}
	fetchURL, err := req.RequireString("url")
	if err != nil {
		return errResult("fetch_url: " + err.Error()), nil
	}

	args := req.GetArguments()
	method := http.MethodGet
	if v, ok := args["method"].(string); ok && v != "" {
		method = strings.ToUpper(v)
	}
	if !fetchURLMethods[method] {
		return errResult(fmt.Sprintf("fetch_url: invalid method %q", method)), nil
	}

	var body string
	if v, ok := args["body"].(string); ok {
		body = v
	}

	var userHeaders map[string]any
	if v, ok := args["headers"].(map[string]any); ok {
		userHeaders = v
	}

	callLog := startToolCall(ctx, d.Logger, FetchURLName, map[string]any{
		"cluster_id": clusterID,
		"url":        fetchURL,
		"method":     method,
	})
	defer callLog.End()

	cl, err := d.ClusterRepo.Get(ctx, clusterID)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_url: cluster " + clusterID + ": " + err.Error()), nil
	}

	allowedHosts := buildAllowedHosts(cl)

	targetURL, err := url.Parse(fetchURL)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_url: parse url: " + err.Error()), nil
	}
	targetHost := strings.ToLower(targetURL.Hostname())
	if !hostAllowed(targetHost, allowedHosts) {
		return errResult(fmt.Sprintf("fetch_url: host %q is not in the cluster allowlist", targetHost)), nil
	}

	if err := httpclient.CheckURL(fetchURL, allowedHosts...); err != nil {
		callLog.WithError(err)
		return errResult("fetch_url: " + err.Error()), nil
	}

	var bodyReader io.Reader
	if body != "" {
		bodyReader = strings.NewReader(body)
	}
	httpReq, err := http.NewRequestWithContext(ctx, method, fetchURL, bodyReader)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_url: build request: " + err.Error()), nil
	}

	for k, v := range userHeaders {
		httpReq.Header.Set(k, fmt.Sprint(v))
	}
	if body != "" {
		httpReq.Header.Set("Content-Length", strconv.Itoa(len(body)))
	}

	// Preserve an explicitly provided Authorization header; otherwise apply cluster auth.
	if httpReq.Header.Get("Authorization") == "" {
		if err := httpclient.ApplyAuth(httpReq, "", cl); err != nil {
			callLog.WithError(err)
			return errResult("fetch_url: auth: " + err.Error()), nil
		}
	}

	resp, err := d.HTTPClient.DoWithRedirect(ctx, httpReq, 5, allowedHosts...)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_url: " + err.Error()), nil
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_url: read body: " + err.Error()), nil
	}

	bodyTruncated := int64(len(respBody)) >= d.MaxResponseBytes &&
		(resp.ContentLength < 0 || resp.ContentLength > d.MaxResponseBytes)

	result := map[string]any{
		"status":         resp.Status,
		"status_code":    resp.StatusCode,
		"headers":        resp.Header,
		"body":           string(respBody),
		"body_truncated": bodyTruncated,
	}
	callLog.WithResult(map[string]any{"status_code": resp.StatusCode, "bytes": len(respBody)})
	return textResult(encodeJSON(result)), nil
}

// buildAllowedHosts extracts hostnames from the cluster's allowlist
// and RM/SHS URLs. Full URLs are parsed and only their host is used; entries
// that do not parse as URLs are passed through as-is so glob patterns such as
// "*.example.com" continue to work.
func buildAllowedHosts(cl *cluster.Cluster) []string {
	var hosts []string
	seen := make(map[string]bool)
	add := func(s string) {
		s = strings.ToLower(strings.TrimSpace(s))
		if s == "" || seen[s] {
			return
		}
		seen[s] = true
		hosts = append(hosts, s)
	}

	for _, raw := range cl.URLAllowlist {
		if u, err := url.Parse(raw); err == nil && u.Hostname() != "" {
			add(u.Hostname())
		} else {
			add(raw)
		}
	}

	for _, raw := range []string{cl.RMURL, cl.SHSURL} {
		if u, err := url.Parse(raw); err == nil && u.Hostname() != "" {
			add(u.Hostname())
		}
	}

	return hosts
}

// hostAllowed reports whether host matches any pattern in allowed. Patterns may
// be exact hostnames, domain suffixes ("example.com" matches "rm.example.com"),
// or wildcard suffixes ("*.example.com"). It mirrors httpclient.matchAllowedHost.
func hostAllowed(host string, allowed []string) bool {
	for _, pattern := range allowed {
		pattern = strings.ToLower(strings.TrimSpace(pattern))
		if pattern == "" {
			continue
		}
		if pattern == host {
			return true
		}
		if strings.HasPrefix(pattern, "*.") {
			pattern = pattern[2:]
		}
		if strings.HasPrefix(pattern, ".") {
			pattern = pattern[1:]
		}
		if host == pattern || strings.HasSuffix(host, "."+pattern) {
			return true
		}
	}
	return false
}