// Copyright 2026 The CLI Printing Press Authors. Licensed under Apache-2.0. See LICENSE.
// Code generated by CLI Printing Press. DO NOT EDIT.
// Generated by CLI Printing Press (https://github.com/mvanhorn/cli-printing-press). DO NOT EDIT.

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

func TestHTTPBindIsLoopback(t *testing.T) {
	t.Parallel()

	cases := []struct {
		addr string
		want bool
	}{
		{addr: "127.0.0.1:7777", want: true},
		{addr: "localhost:7777", want: true},
		{addr: "[::1]:7777", want: true},
		{addr: ":7777", want: false},
		{addr: "0.0.0.0:7777", want: false},
		{addr: "[::]:7777", want: false},
		{addr: "192.168.1.10:7777", want: false},
		{addr: "127.0.0.1", want: false},
		{addr: "example.invalid:7777", want: false},
	}
	for _, tc := range cases {
		if got := httpBindIsLoopback(tc.addr); got != tc.want {
			t.Fatalf("httpBindIsLoopback(%q) = %v, want %v", tc.addr, got, tc.want)
		}
	}
}

func TestClassifyHTTPBindPinsNamedLoopback(t *testing.T) {
	t.Parallel()

	pinned, ok := classifyHTTPBind("localhost:7777")
	if !ok {
		t.Fatal("localhost:7777 must classify as loopback")
	}
	host, port, err := net.SplitHostPort(pinned)
	if err != nil {
		t.Fatalf("pinned addr %q: %v", pinned, err)
	}
	ip := net.ParseIP(host)
	if ip == nil || !ip.IsLoopback() {
		t.Fatalf("classifyHTTPBind(%q) = %q, want a loopback IP", "localhost:7777", pinned)
	}
	if port != "7777" {
		t.Fatalf("pinned port = %q, want 7777", port)
	}

	same, ipLoopback := classifyHTTPBind("127.0.0.1:7777")
	if !ipLoopback || same != "127.0.0.1:7777" {
		t.Fatalf("classifyHTTPBind(127.0.0.1:7777) = %q, %v", same, ipLoopback)
	}
	open, openLoopback := classifyHTTPBind("0.0.0.0:7777")
	if openLoopback || open != "0.0.0.0:7777" {
		t.Fatalf("classifyHTTPBind(0.0.0.0:7777) = %q, %v", open, openLoopback)
	}
}

func TestRequireTLSForNonLoopback(t *testing.T) {
	t.Parallel()

	if err := requireTLSForNonLoopback("127.0.0.1:7777", true, "", ""); err != nil {
		t.Fatalf("loopback without TLS: %v", err)
	}
	if err := requireTLSForNonLoopback("0.0.0.0:7777", false, "", ""); err == nil {
		t.Fatal("non-loopback without TLS must fail")
	}
	if err := requireTLSForNonLoopback(":7777", false, "cert.pem", ""); err == nil {
		t.Fatal("empty-host bind without both TLS files must fail")
	}
	if err := requireTLSForNonLoopback("0.0.0.0:7777", false, "cert.pem", "key.pem"); err != nil {
		t.Fatalf("non-loopback with TLS files: %v", err)
	}
}

func TestRequireHTTPCallerToken(t *testing.T) {
	t.Setenv(httpTokenEnvVar, "")
	if _, err := requireHTTPCallerToken(); err == nil {
		t.Fatal("empty token env must fail")
	}
	t.Setenv(httpTokenEnvVar, "   ")
	if _, err := requireHTTPCallerToken(); err == nil {
		t.Fatal("whitespace token env must fail")
	}
	t.Setenv(httpTokenEnvVar, "shared-secret")
	got, err := requireHTTPCallerToken()
	if err != nil {
		t.Fatalf("set token env: %v", err)
	}
	if got != "shared-secret" {
		t.Fatalf("requireHTTPCallerToken() = %q, want shared-secret", got)
	}
}

func TestRequireBearerAuthRejectsMissingAndWrongTokens(t *testing.T) {
	t.Parallel()

	inner := server.NewStreamableHTTPServer(server.NewMCPServer("auth-proof", "0.0.0-dev"))
	handler := requireBearerAuth("shared-secret", inner)
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	resp, err := http.Post(srv.URL+"/mcp", "application/json", strings.NewReader(`{}`))
	if err != nil {
		t.Fatalf("unauthenticated POST: %v", err)
	}
	body, _ := io.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("unauthenticated POST status = %d, want 401; body %s", resp.StatusCode, body)
	}

	req, err := http.NewRequest(http.MethodPost, srv.URL+"/mcp", strings.NewReader(`{}`))
	if err != nil {
		t.Fatalf("new request: %v", err)
	}
	req.Header.Set("Authorization", "Bearer wrong-secret")
	req.Header.Set("Content-Type", "application/json")
	resp, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("wrong-token POST: %v", err)
	}
	body, _ = io.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("wrong-token POST status = %d, want 401; body %s", resp.StatusCode, body)
	}

	req, err = http.NewRequest(http.MethodPost, srv.URL+"/mcp", strings.NewReader(`{}`))
	if err != nil {
		t.Fatalf("new authed request: %v", err)
	}
	req.Header.Set("Authorization", "Bearer shared-secret")
	req.Header.Set("Content-Type", "application/json")
	resp, err = http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("authed POST: %v", err)
	}
	body, _ = io.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode == http.StatusUnauthorized {
		t.Fatalf("matching bearer still unauthorized; body %s", body)
	}
}

func TestMCPHTTPServerLimits(t *testing.T) {
	t.Parallel()

	if mcpHTTPMaxBodyBytes != 4<<20 {
		t.Fatalf("max body = %d, want %d", mcpHTTPMaxBodyBytes, 4<<20)
	}
	if mcpHTTPBodyReadDeadline != 30*time.Second {
		t.Fatalf("body read deadline = %s, want 30s", mcpHTTPBodyReadDeadline)
	}
	srv := newMCPHTTPServer("127.0.0.1:0", http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
	if srv.IdleTimeout != 2*time.Minute {
		t.Fatalf("IdleTimeout = %s, want 2m", srv.IdleTimeout)
	}
	if srv.ReadHeaderTimeout != 10*time.Second {
		t.Fatalf("ReadHeaderTimeout = %s, want 10s", srv.ReadHeaderTimeout)
	}
	if srv.ReadTimeout != 0 || srv.WriteTimeout != 0 {
		t.Fatalf("server-wide deadlines must stay unset, read=%s write=%s", srv.ReadTimeout, srv.WriteTimeout)
	}
}

func TestMCPHTTPDeclaredOversizeBodyIsRefused(t *testing.T) {
	t.Parallel()

	var called atomic.Bool
	handler := requireBearerAuth("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called.Store(true)
		w.WriteHeader(http.StatusOK)
	}))
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	status, err := mcpHTTPRawStatus(t, srv, "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer secret\r\nContent-Type: application/json\r\nContent-Length: "+strconv.FormatInt(mcpHTTPMaxBodyBytes+1, 10)+"\r\n\r\n")
	if err != nil {
		t.Fatalf("read status: %v", err)
	}
	if status != http.StatusRequestEntityTooLarge {
		t.Fatalf("status = %d, want 413", status)
	}
	if called.Load() {
		t.Fatal("handler ran for a declared oversize body")
	}
}

func TestMCPHTTPChunkedOversizeBodyIsRefused(t *testing.T) {
	t.Parallel()

	var called atomic.Bool
	handler := requireBearerAuthWithLimits("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called.Store(true)
		w.WriteHeader(http.StatusOK)
	}), 32, 5*time.Second)
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	body := bytes.Repeat([]byte("x"), 64)
	req, err := http.NewRequest(http.MethodPost, srv.URL+"/mcp", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.ContentLength = -1
	req.Header.Del("Content-Length")
	req.Header.Set("Authorization", "Bearer secret")
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("chunked POST: %v", err)
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)
	if resp.StatusCode != http.StatusRequestEntityTooLarge {
		t.Fatalf("status = %d, want 413", resp.StatusCode)
	}
	if called.Load() {
		t.Fatal("handler ran for a chunked oversize body")
	}
}

func TestMCPHTTPTrickledBodyTimesOut(t *testing.T) {
	t.Parallel()

	var called atomic.Bool
	handler := requireBearerAuthWithLimits("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called.Store(true)
		w.WriteHeader(http.StatusOK)
	}), 1<<20, 200*time.Millisecond)
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	pr, pw := io.Pipe()
	go func() {
		_, _ = pw.Write([]byte(`{"jsonrpc":"2.0"}`))
		time.Sleep(time.Second)
		_ = pw.Close()
	}()
	req, err := http.NewRequest(http.MethodPost, srv.URL+"/mcp", pr)
	if err != nil {
		t.Fatal(err)
	}
	req.ContentLength = -1
	req.Header.Set("Authorization", "Bearer secret")
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{Timeout: 3 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatalf("trickled POST: %v", err)
	}
	defer resp.Body.Close()
	payload, readErr := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusRequestTimeout {
		t.Fatalf("status = %d, want 408; body %s; read %v", resp.StatusCode, payload, readErr)
	}
	if called.Load() {
		t.Fatal("handler ran for a stalled body")
	}
	if _, err := resp.Body.Read(make([]byte, 1)); err != io.EOF {
		t.Fatalf("connection after 408 = %v, want EOF", err)
	}
}

func TestMCPHTTPUnauthenticatedSkipsBodyRead(t *testing.T) {
	t.Parallel()

	var called atomic.Bool
	handler := requireBearerAuth("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called.Store(true)
		w.WriteHeader(http.StatusOK)
	}))
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	start := time.Now()
	status, err := mcpHTTPRawStatus(t, srv, "POST /mcp HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Type: application/json\r\nContent-Length: 1000000\r\n\r\n")
	if err != nil {
		t.Fatalf("read status: %v", err)
	}
	if time.Since(start) > 800*time.Millisecond {
		t.Fatalf("401 took %s; body read should not block the refusal", time.Since(start))
	}
	if status != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401", status)
	}
	if called.Load() {
		t.Fatal("handler ran before authentication")
	}
}

func TestMCPHTTPSmallPostReachesHandler(t *testing.T) {
	t.Parallel()

	var got []byte
	handler := requireBearerAuthWithLimits("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got, _ = io.ReadAll(r.Body)
		w.WriteHeader(http.StatusNoContent)
	}), 128, 200*time.Millisecond)
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	req, err := http.NewRequest(http.MethodPost, srv.URL+"/mcp", strings.NewReader("ping"))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer secret")
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusNoContent {
		t.Fatalf("status = %d, want 204", resp.StatusCode)
	}
	if string(got) != "ping" {
		t.Fatalf("handler body = %q, want ping", got)
	}
}

func TestMCPHTTPGETSurvivesBodyReadDeadline(t *testing.T) {
	t.Parallel()
	assertMCPHTTPStreamSurvivesDeadline(t, false, http.MethodGet)
}

func TestMCPHTTPPostResponseSurvivesBodyReadDeadline(t *testing.T) {
	t.Parallel()
	assertMCPHTTPStreamSurvivesDeadline(t, false, http.MethodPost)
}

func TestMCPHTTPPostResponseSurvivesBodyReadDeadlineHTTP2(t *testing.T) {
	t.Parallel()
	assertMCPHTTPStreamSurvivesDeadline(t, true, http.MethodPost)
}

func TestMCPHTTPIdleKeepAliveIsReaped(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	handler := requireBearerAuth("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, "ok\n")
	}))
	srv := newMCPHTTPServer(ln.Addr().String(), handler)
	srv.IdleTimeout = 200 * time.Millisecond
	go func() { _ = srv.Serve(ln) }()
	t.Cleanup(func() { _ = srv.Close() })

	conn, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	if _, err := io.WriteString(conn, "GET /mcp HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer secret\r\nConnection: keep-alive\r\n\r\n"); err != nil {
		t.Fatal(err)
	}
	br := bufio.NewReader(conn)
	resp, err := http.ReadResponse(br, nil)
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.Copy(io.Discard, resp.Body)
	_ = resp.Body.Close()
	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	if _, err := br.Read(make([]byte, 1)); err == nil {
		t.Fatal("idle keep-alive stayed open past IdleTimeout")
	}
}

func TestMCPHTTPInitializeIsNotLimited(t *testing.T) {
	t.Parallel()

	inner := server.NewStreamableHTTPServer(server.NewMCPServer("limit-proof", "0.0.0-dev"))
	handler := requireBearerAuth("secret", inner)
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"0"}}}`
	req, err := http.NewRequest(http.MethodPost, srv.URL+"/mcp", strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer secret")
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	payload, _ := io.ReadAll(resp.Body)
	switch resp.StatusCode {
	case http.StatusUnauthorized, http.StatusRequestEntityTooLarge, http.StatusRequestTimeout:
		t.Fatalf("initialize status = %d; body %s", resp.StatusCode, payload)
	}
}

func assertMCPHTTPStreamSurvivesDeadline(t *testing.T, http2 bool, method string) {
	t.Helper()

	readTimeout := 200 * time.Millisecond
	var sawBody string
	handler := requireBearerAuthWithLimits("secret", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if method == http.MethodPost {
			buf, err := io.ReadAll(r.Body)
			if err != nil {
				t.Errorf("handler read body: %v", err)
				return
			}
			sawBody = string(buf)
		}
		_, _ = io.WriteString(w, "a")
		if flusher, ok := w.(http.Flusher); ok {
			flusher.Flush()
		}
		time.Sleep(readTimeout + 150*time.Millisecond)
		_, _ = io.WriteString(w, "b")
	}), 1024, readTimeout)

	var srv *httptest.Server
	if http2 {
		srv = httptest.NewUnstartedServer(handler)
		srv.EnableHTTP2 = true
		srv.StartTLS()
	} else {
		srv = httptest.NewServer(handler)
	}
	t.Cleanup(srv.Close)

	var body io.Reader
	if method == http.MethodPost {
		body = strings.NewReader("ping")
	}
	req, err := http.NewRequest(method, srv.URL+"/mcp", body)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Authorization", "Bearer secret")
	if method == http.MethodPost {
		req.Header.Set("Content-Type", "application/json")
	}
	client := srv.Client()
	client.Timeout = 3 * time.Second
	resp, err := client.Do(req)
	if err != nil {
		t.Fatalf("%s http2=%v: %v", method, http2, err)
	}
	defer resp.Body.Close()
	payload, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read %s body: %v", method, err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("%s status = %d; body %s", method, resp.StatusCode, payload)
	}
	if string(payload) != "ab" {
		t.Fatalf("%s body = %q, want ab", method, payload)
	}
	if method == http.MethodPost && sawBody != "ping" {
		t.Fatalf("handler saw %q, want ping", sawBody)
	}
}

func mcpHTTPRawStatus(t *testing.T, srv *httptest.Server, request string) (int, error) {
	t.Helper()
	conn, err := net.Dial("tcp", srv.Listener.Addr().String())
	if err != nil {
		return 0, err
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
	if _, err := io.WriteString(conn, request); err != nil {
		return 0, err
	}
	br := bufio.NewReader(conn)
	line, err := br.ReadString('\n')
	if err != nil {
		return 0, err
	}
	var status int
	if _, err := fmt.Sscanf(line, "HTTP/1.1 %d", &status); err != nil {
		return 0, err
	}
	return status, nil
}
