package server

import (
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestRateLimiter(t *testing.T) {
	now := time.Now()
	limiter := NewRateLimiter(2, time.Hour)
	limiter.now = func() time.Time { return now }

	if ok, _ := limiter.Allow("a"); !ok {
		t.Fatal("first event should pass")
	}
	if ok, _ := limiter.Allow("a"); !ok {
		t.Fatal("second event should pass")
	}
	ok, retry := limiter.Allow("a")
	if ok || retry <= 0 || retry > time.Hour {
		t.Fatalf("third event: ok=%v retry=%v", ok, retry)
	}
	if ok, _ := limiter.Allow("b"); !ok {
		t.Fatal("another key should have its own window")
	}

	now = now.Add(time.Hour + time.Second)
	if ok, _ := limiter.Allow("a"); !ok {
		t.Fatal("a new window should allow events again")
	}
}

func TestClientIP(t *testing.T) {
	cases := []struct {
		forwarded, real, remote, want string
	}{
		{"203.0.113.9, 10.0.0.1", "", "10.0.0.2:1234", "203.0.113.9"},
		{"", "198.51.100.7", "10.0.0.2:1234", "198.51.100.7"},
		{"", "", "192.0.2.5:443", "192.0.2.5"},
		{"garbage", "", "192.0.2.5:443", "192.0.2.5"},
	}
	for _, c := range cases {
		r := httptest.NewRequest("GET", "/", nil)
		r.RemoteAddr = c.remote
		if c.forwarded != "" {
			r.Header.Set("X-Forwarded-For", c.forwarded)
		}
		if c.real != "" {
			r.Header.Set("X-Real-IP", c.real)
		}
		if got := clientIP(r); got != c.want {
			t.Errorf("clientIP(%q, %q, %q) = %q, want %q", c.forwarded, c.real, c.remote, got, c.want)
		}
	}
}

func TestCreateRateLimit(t *testing.T) {
	store, err := OpenStore("")
	if err != nil {
		t.Fatal(err)
	}
	api := &API{
		Store: store, Hub: NewHub(store), Domain: "opentunnel.test", StartedAt: time.Now(),
		CreateLimiter: NewRateLimiter(1, time.Hour),
	}
	srv := httptest.NewServer(api)
	defer srv.Close()

	resp := request(t, "POST", srv.URL+"/api/tunnel", "", "application/json", `{}`)
	if resp.StatusCode != 201 {
		t.Fatalf("first create: %d", resp.StatusCode)
	}
	resp = request(t, "POST", srv.URL+"/api/tunnel", "", "application/json", `{}`)
	body := bodyString(t, resp)
	if resp.StatusCode != 429 || !strings.Contains(body, `"_tag":"RateLimitedError"`) {
		t.Fatalf("second create: %d %q", resp.StatusCode, body)
	}
	if resp.Header.Get("Retry-After") == "" {
		t.Error("429 response is missing Retry-After")
	}
}