ollama/server/cloud_proxy_test.go

package server

import (
	"net/http"
	"testing"

	"github.com/gin-gonic/gin"
)

func TestCopyProxyRequestHeaders_StripsConnectionTokenHeaders(t *testing.T) {
	src := http.Header{}
	src.Add("Connection", "keep-alive, X-Trace-Hop, x-alt-hop")
	src.Add("X-Trace-Hop", "drop-me")
	src.Add("X-Alt-Hop", "drop-me-too")
	src.Add("Keep-Alive", "timeout=5")
	src.Add("X-End-To-End", "keep-me")

	dst := http.Header{}
	copyProxyRequestHeaders(dst, src)

	if got := dst.Get("Connection"); got != "" {
		t.Fatalf("expected Connection to be stripped, got %q", got)
	}
	if got := dst.Get("Keep-Alive"); got != "" {
		t.Fatalf("expected Keep-Alive to be stripped, got %q", got)
	}
	if got := dst.Get("X-Trace-Hop"); got != "" {
		t.Fatalf("expected X-Trace-Hop to be stripped via Connection token, got %q", got)
	}
	if got := dst.Get("X-Alt-Hop"); got != "" {
		t.Fatalf("expected X-Alt-Hop to be stripped via Connection token, got %q", got)
	}
	if got := dst.Get("X-End-To-End"); got != "keep-me" {
		t.Fatalf("expected X-End-To-End to be forwarded, got %q", got)
	}
}

func TestCopyProxyResponseHeaders_StripsConnectionTokenHeaders(t *testing.T) {
	src := http.Header{}
	src.Add("Connection", "X-Upstream-Hop")
	src.Add("X-Upstream-Hop", "drop-me")
	src.Add("Content-Type", "application/json")
	src.Add("X-Server-Trace", "keep-me")

	dst := http.Header{}
	copyProxyResponseHeaders(dst, src)

	if got := dst.Get("Connection"); got != "" {
		t.Fatalf("expected Connection to be stripped, got %q", got)
	}
	if got := dst.Get("X-Upstream-Hop"); got != "" {
		t.Fatalf("expected X-Upstream-Hop to be stripped via Connection token, got %q", got)
	}
	if got := dst.Get("Content-Type"); got != "application/json" {
		t.Fatalf("expected Content-Type to be forwarded, got %q", got)
	}
	if got := dst.Get("X-Server-Trace"); got != "keep-me" {
		t.Fatalf("expected X-Server-Trace to be forwarded, got %q", got)
	}
}

func TestResolveCloudProxyBaseURL_Default(t *testing.T) {
	baseURL, signingHost, overridden, err := resolveCloudProxyBaseURL("", gin.ReleaseMode)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if overridden {
		t.Fatal("expected override=false for empty input")
	}
	if baseURL != defaultCloudProxyBaseURL {
		t.Fatalf("expected default base URL %q, got %q", defaultCloudProxyBaseURL, baseURL)
	}
	if signingHost != defaultCloudProxySigningHost {
		t.Fatalf("expected default signing host %q, got %q", defaultCloudProxySigningHost, signingHost)
	}
}

func TestResolveCloudProxyBaseURL_ReleaseAllowsLoopback(t *testing.T) {
	baseURL, signingHost, overridden, err := resolveCloudProxyBaseURL("http://localhost:8080", gin.ReleaseMode)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !overridden {
		t.Fatal("expected override=true")
	}
	if baseURL != "http://localhost:8080" {
		t.Fatalf("unexpected base URL: %q", baseURL)
	}
	if signingHost != "localhost" {
		t.Fatalf("unexpected signing host: %q", signingHost)
	}
}

func TestResolveCloudProxyBaseURL_ReleaseRejectsNonLoopback(t *testing.T) {
	_, _, _, err := resolveCloudProxyBaseURL("https://example.com", gin.ReleaseMode)
	if err == nil {
		t.Fatal("expected error for non-loopback override in release mode")
	}
}

func TestResolveCloudProxyBaseURL_DevAllowsNonLoopbackHTTPS(t *testing.T) {
	baseURL, signingHost, overridden, err := resolveCloudProxyBaseURL("https://example.com:8443", gin.DebugMode)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !overridden {
		t.Fatal("expected override=true")
	}
	if baseURL != "https://example.com:8443" {
		t.Fatalf("unexpected base URL: %q", baseURL)
	}
	if signingHost != "example.com" {
		t.Fatalf("unexpected signing host: %q", signingHost)
	}
}

func TestResolveCloudProxyBaseURL_DevRejectsNonLoopbackHTTP(t *testing.T) {
	_, _, _, err := resolveCloudProxyBaseURL("http://example.com", gin.DebugMode)
	if err == nil {
		t.Fatal("expected error for non-loopback http override in dev mode")
	}
}

func TestBuildCloudSignatureChallengeIncludesExistingQuery(t *testing.T) {
	req, err := http.NewRequest(http.MethodPost, "https://ollama.com/v1/messages?beta=true&foo=bar", nil)
	if err != nil {
		t.Fatalf("failed to create request: %v", err)
	}

	got := buildCloudSignatureChallenge(req, "123")
	want := "POST,/v1/messages?beta=true&foo=bar&ts=123"
	if got != want {
		t.Fatalf("challenge mismatch: got %q want %q", got, want)
	}
	if req.URL.RawQuery != "beta=true&foo=bar&ts=123" {
		t.Fatalf("unexpected signed query: %q", req.URL.RawQuery)
	}
}

func TestBuildCloudSignatureChallengeOverwritesExistingTimestamp(t *testing.T) {
	req, err := http.NewRequest(http.MethodPost, "https://ollama.com/v1/messages?beta=true&ts=999", nil)
	if err != nil {
		t.Fatalf("failed to create request: %v", err)
	}

	got := buildCloudSignatureChallenge(req, "123")
	want := "POST,/v1/messages?beta=true&ts=123"
	if got != want {
		t.Fatalf("challenge mismatch: got %q want %q", got, want)
	}
	if req.URL.RawQuery != "beta=true&ts=123" {
		t.Fatalf("unexpected signed query: %q", req.URL.RawQuery)
	}
}