package server

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

func guard(auth bool, hosts ...string) RequestGuard {
	return RequestGuard{
		AuthOn: auth, AllowedHosts: hosts,
		Next: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusNoContent)
		}),
	}
}

func TestRequestGuardWrites(t *testing.T) {
	cases := []struct {
		name   string
		auth   bool
		host   string
		origin string
		site   string
		want   int
	}{
		{"same origin", false, "127.0.0.1:8788", "http://127.0.0.1:8788", "", http.StatusNoContent},
		{"cross origin", false, "127.0.0.1:8788", "https://evil.example", "", http.StatusForbidden},
		{"null origin", false, "127.0.0.1:8788", "null", "", http.StatusForbidden},
		{"sec-fetch cross", false, "127.0.0.1:8788", "", "cross-site", http.StatusForbidden},
		{"sec-fetch same", false, "127.0.0.1:8788", "", "same-origin", http.StatusNoContent},
		{"no browser headers", false, "127.0.0.1:8788", "", "", http.StatusNoContent},
		{"allowed host origin", false, "127.0.0.1:8788", "https://app.example", "", http.StatusNoContent},
		{"evil host without auth", false, "evil.example", "", "", http.StatusForbidden},
		{"evil host with auth", true, "evil.example", "", "", http.StatusNoContent},
		{"ip host", false, "192.168.1.9:8788", "", "", http.StatusNoContent},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, "/api/memory", strings.NewReader(`{}`))
			req.Host = tc.host
			if tc.origin != "" {
				req.Header.Set("Origin", tc.origin)
			}
			if tc.site != "" {
				req.Header.Set("Sec-Fetch-Site", tc.site)
			}
			rec := httptest.NewRecorder()
			hosts := []string{}
			if tc.name == "allowed host origin" {
				hosts = []string{"app.example"}
			}
			guard(tc.auth, hosts...).ServeHTTP(rec, req)
			if rec.Code != tc.want {
				t.Fatalf("status %d want %d body %s", rec.Code, tc.want, rec.Body.String())
			}
			if rec.Code == http.StatusForbidden && !strings.Contains(rec.Body.String(), "allowedHosts") {
				t.Fatalf("refusal should mention allowedHosts: %s", rec.Body.String())
			}
		})
	}
}

func TestPageAndArtifactHeaders(t *testing.T) {
	rec := httptest.NewRecorder()
	setPageHeaders(rec)
	if rec.Header().Get("Content-Security-Policy") == "" || rec.Header().Get("X-Content-Type-Options") != "nosniff" || rec.Header().Get("Referrer-Policy") != "same-origin" {
		t.Fatalf("page headers: %v", rec.Header())
	}
	art := httptest.NewRecorder()
	setArtifactHeaders(art, "text/html; charset=utf-8")
	if !strings.Contains(art.Header().Get("Content-Security-Policy"), "sandbox") || art.Header().Get("X-Content-Type-Options") != "nosniff" {
		t.Fatalf("artifact headers: %v", art.Header())
	}
	plain := httptest.NewRecorder()
	setArtifactHeaders(plain, "text/plain")
	if plain.Header().Get("Content-Security-Policy") != "" {
		t.Fatalf("plain artifact should not be sandboxed: %v", plain.Header())
	}
}