Local-Cursor/cursor-tab-server/main_test.go

package main

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestAuthorizationAndUpstreamTokenSeparation(t *testing.T) {
	calls := 0
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls++
		if r.Header.Get("Authorization") != "Bearer upstream-secret" {
			t.Error("wrong upstream token")
		}
		_, _ = w.Write([]byte("ok"))
	}))
	defer upstream.Close()
	app := newServerApp(appConfig{Token: "upstream-secret", ClientToken: "client-secret"}, upstream.Client(), map[string]string{"/rpc": upstream.URL})
	for _, token := range []string{"", "Bearer wrong", "Bearer upstream-secret", "Bearer client-secret"} {
		r := httptest.NewRequest(http.MethodPost, "/rpc", strings.NewReader("request"))
		r.Header.Set("Authorization", token)
		w := httptest.NewRecorder()
		app.ServeHTTP(w, r)
		want := 401
		if token == "Bearer client-secret" {
			want = 200
		}
		if w.Code != want {
			t.Fatalf("status=%d want=%d", w.Code, want)
		}
	}
	if calls != 1 {
		t.Fatal("unauthorized request reached upstream")
	}
}

func TestOversizeRejectedBeforeUpstream(t *testing.T) {
	app := newServerApp(appConfig{Token: "upstream", ClientToken: "client"}, nil, map[string]string{"/rpc": "http://127.0.0.1:1"})
	r := httptest.NewRequest(http.MethodPost, "/rpc", strings.NewReader(""))
	r.ContentLength = maxRequestBytes + 1
	r.Header.Set("Authorization", "Bearer client")
	w := httptest.NewRecorder()
	app.ServeHTTP(w, r)
	if w.Code != 413 {
		t.Fatalf("status=%d", w.Code)
	}
}

type testTransport func(*http.Request) (*http.Response, error)

func (t testTransport) RoundTrip(r *http.Request) (*http.Response, error) { return t(r) }

type brokenStream struct{ sent bool }

func (b *brokenStream) Read(p []byte) (int, error) {
	if !b.sent {
		b.sent = true
		return copy(p, []byte("partial-event")), nil
	}
	return 0, errors.New("stream interrupted")
}
func (b *brokenStream) Close() error { return nil }

func TestStreamErrorNeverAppendsHTTPError(t *testing.T) {
	client := &http.Client{Transport: testTransport(func(*http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: 200, Header: http.Header{}, Body: &brokenStream{}}, nil
	})}
	app := newServerApp(appConfig{Token: "upstream", ClientToken: "client"}, client, map[string]string{"/rpc": "https://upstream.invalid"})
	r := httptest.NewRequest(http.MethodPost, "/rpc", nil)
	r.Header.Set("Authorization", "Bearer client")
	w := httptest.NewRecorder()
	app.ServeHTTP(w, r)
	if w.Code != 200 || w.Body.String() != "partial-event" {
		t.Fatalf("response corrupted: %d %q", w.Code, w.Body.String())
	}
}

func TestCanceledRequestDoesNotCallUpstream(t *testing.T) {
	client := &http.Client{Transport: testTransport(func(r *http.Request) (*http.Response, error) { return nil, r.Context().Err() })}
	app := newServerApp(appConfig{Token: "upstream", ClientToken: "client"}, client, map[string]string{"/rpc": "https://upstream.invalid"})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	r := httptest.NewRequest(http.MethodPost, "/rpc", nil).WithContext(ctx)
	r.Header.Set("Authorization", "Bearer client")
	w := httptest.NewRecorder()
	app.ServeHTTP(w, r)
	if w.Code != 499 {
		t.Fatalf("status=%d", w.Code)
	}
}

func TestClientTokenRequiredAndDistinct(t *testing.T) {
	p := filepath.Join(t.TempDir(), "config.yaml")
	for _, content := range []string{"token: upstream\n", "token: same\nclientToken: same\n"} {
		if err := os.WriteFile(p, []byte(content), 0600); err != nil {
			t.Fatal(err)
		}
		if _, err := loadConfig(p); err == nil {
			t.Fatal("insecure configuration accepted")
		}
	}
	_ = os.WriteFile(p, []byte("token: upstream\nclientToken: client\n"), 0600)
	if _, err := loadConfig(p); err != nil {
		t.Fatal(err)
	}
}

func TestUpstreamRedirectDoesNotForwardSecret(t *testing.T) {
	leaked := false
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { leaked = true; _, _ = io.WriteString(w, "bad") }))
	defer target.Close()
	redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
	defer redirect.Close()
	r, _ := http.NewRequest(http.MethodPost, redirect.URL, strings.NewReader(""))
	r.Header.Set("Authorization", "Bearer secret")
	resp, err := newHTTPClient().Do(r)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 307 || leaked {
		t.Fatal("redirect followed with credentials")
	}
}