package main

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestSpoofedForwardedRemoved(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("X-Forwarded-For") != "192.0.2.1" || r.Header.Get("Forwarded") != "" || r.Header.Get("X-Forwarded-Proto") != "http" {
			t.Errorf("unexpected headers: %v", r.Header)
		}
		w.WriteHeader(204)
	}))
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)
	p := proxy(u, time.Second)
	defer p.Transport.(*http.Transport).CloseIdleConnections()
	r := httptest.NewRequest("GET", "http://example.test/", nil)
	r.RemoteAddr = "192.0.2.1:42"
	r.Header.Set("X-Forwarded-For", "198.51.100.99")
	r.Header.Set("Forwarded", "for=198.51.100.99")
	r.Header.Set("X-Forwarded-Proto", "https")
	w := httptest.NewRecorder()
	p.ServeHTTP(w, r)
	if w.Code != 204 {
		t.Fatal(w.Code)
	}
}
func TestHeaderTimeout(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
	defer upstream.Close()
	u, _ := url.Parse(upstream.URL)
	p := proxy(u, 20*time.Millisecond)
	defer p.Transport.(*http.Transport).CloseIdleConnections()
	w := httptest.NewRecorder()
	p.ServeHTTP(w, httptest.NewRequest("GET", "http://example.test/", nil))
	if w.Code != 502 || !strings.Contains(w.Body.String(), "bad gateway") {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
}
func TestBucketConcurrent(t *testing.T) {
	now := time.Unix(0, 0)
	b := bucket{tokens: 2, capacity: 2, rate: 1, last: now}
	var accepted atomic.Int32
	var wg sync.WaitGroup
	for range 20 {
		wg.Go(func() {
			if b.allow(now) {
				accepted.Add(1)
			}
		})
	}
	wg.Wait()
	if accepted.Load() != 2 {
		t.Fatal(accepted.Load())
	}
	if !b.allow(now.Add(time.Second)) {
		t.Fatal("expected refill")
	}
}