package cors

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"
)

func exampleConfig() Config {
	return Config{Enabled: true, AllowedOrigins: []string{"https://admin.example", "https://merchant.example"}, AllowedMethods: []string{"GET", "POST", "DELETE"}, AllowedHeaders: []string{"Authorization", "Content-Type"}, ExposedHeaders: []string{"X-Request-ID", "Content-Disposition"}, MaxAge: 10 * time.Minute}
}

func TestPolicyPreflightAndActual(t *testing.T) {
	for _, tc := range []struct {
		name, method, origin, requestedMethod, requestedHeaders string
		status                                                  int
		called, allowed                                         bool
	}{
		{"preflight", "OPTIONS", "https://admin.example", "POST", "authorization, CONTENT-TYPE", 204, false, true},
		{"no request headers", "OPTIONS", "https://admin.example", "DELETE", "", 204, false, true},
		{"unknown origin", "OPTIONS", "https://other.example", "POST", "authorization", 403, false, false},
		{"suffix attack", "OPTIONS", "https://admin.example.evil", "POST", "", 403, false, false},
		{"scheme mismatch", "OPTIONS", "http://admin.example", "POST", "", 403, false, false},
		{"port mismatch", "OPTIONS", "https://admin.example:444", "POST", "", 403, false, false},
		{"opaque origin", "OPTIONS", "null", "POST", "", 403, false, false},
		{"unknown method", "OPTIONS", "https://admin.example", "PATCH", "", 403, false, false},
		{"case sensitive method", "OPTIONS", "https://admin.example", "post", "", 403, false, false},
		{"unknown header", "OPTIONS", "https://admin.example", "POST", "X-Internal", 403, false, false},
		{"trailing empty header", "OPTIONS", "https://admin.example", "POST", "Authorization,", 403, false, false},
		{"control header", "OPTIONS", "https://admin.example", "POST", "Authorization\r\nX-Evil", 403, false, false},
		{"duplicate allowed header", "OPTIONS", "https://admin.example", "POST", "Authorization, authorization", 204, false, true},
		{"oversized list", "OPTIONS", "https://admin.example", "POST", strings.Repeat("Authorization,", 129) + "Authorization", 403, false, false},
		{"actual error response", "GET", "https://merchant.example", "", "", 401, true, true},
		{"actual denied origin still authenticates", "POST", "https://other.example", "", "", 401, true, false},
		{"actual denied method still authenticates", "PUT", "https://admin.example", "", "", 401, true, false},
		{"server callback", "POST", "", "", "", 401, true, false},
		{"ordinary OPTIONS", "OPTIONS", "https://admin.example", "", "", 401, true, false},
		{"OPTIONS without Origin", "OPTIONS", "", "POST", "", 401, true, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			policy, err := New(exampleConfig())
			if err != nil {
				t.Fatal(err)
			}
			called := false
			h := policy.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				called = true
				w.WriteHeader(401)
				io.WriteString(w, "authentication required")
			}))
			r := httptest.NewRequest(tc.method, "http://api.example/resource", nil)
			if tc.origin != "" {
				r.Header.Set("Origin", tc.origin)
			}
			if tc.requestedMethod != "" {
				r.Header.Set("Access-Control-Request-Method", tc.requestedMethod)
			}
			if tc.requestedHeaders != "" {
				r.Header.Set("Access-Control-Request-Headers", tc.requestedHeaders)
			}
			w := httptest.NewRecorder()
			w.Header().Add("Vary", "Accept-Encoding")
			h.ServeHTTP(w, r)
			if w.Code != tc.status || called != tc.called {
				t.Fatalf("status=%d called=%v", w.Code, called)
			}
			if got := w.Header().Get("Access-Control-Allow-Origin"); (got != "") != tc.allowed || tc.allowed && got != tc.origin {
				t.Fatal(w.Header())
			}
			if !strings.Contains(strings.Join(w.Header().Values("Vary"), ","), "Origin") || !strings.Contains(strings.Join(w.Header().Values("Vary"), ","), "Accept-Encoding") {
				t.Fatal(w.Header())
			}
			if !tc.called {
				if w.Body.Len() != 0 {
					t.Fatal("preflight wrapped in business response")
				}
				if !strings.Contains(strings.Join(w.Header().Values("Vary"), ","), "Access-Control-Request-Headers") {
					t.Fatal(w.Header())
				}
				if tc.allowed && (w.Header().Get("Access-Control-Max-Age") != "600" || !strings.Contains(w.Header().Get("Access-Control-Allow-Headers"), "Authorization")) {
					t.Fatal(w.Header())
				}
			} else if w.Body.String() != "authentication required" {
				t.Fatal(w.Body.String())
			}
			if tc.called && tc.allowed && w.Header().Get("Access-Control-Expose-Headers") != "X-Request-Id, Content-Disposition" {
				t.Fatal(w.Header())
			}
		})
	}
}

func TestPolicyValidation(t *testing.T) {
	for _, origin := range []string{"https://admin.example/", "https://user:password@admin.example", "null", "https://*.example", "https://admin.example?secret", "https://admin.example#secret", "file://host", "https://host:", "https://host:0", "https://host:65536", "https://host%2eevil", "https://host\r\nX: value", "https://例子.example", "https://[localhost]", "https://[::1%25zone]"} {
		c := exampleConfig()
		c.AllowedOrigins = []string{origin}
		err := c.Validate()
		if err == nil || strings.Contains(err.Error(), origin) {
			t.Fatalf("invalid origin accepted or exposed: %q %v", origin, err)
		}
	}
	for _, change := range []func(*Config){
		func(c *Config) { c.AllowedOrigins = nil },
		func(c *Config) { c.AllowedOrigins = []string{"*"}; c.AllowCredentials = true },
		func(c *Config) { c.AllowedOrigins = []string{"*", "https://admin.example"} },
		func(c *Config) { c.AllowedHeaders = []string{"*"} },
		func(c *Config) { c.ExposedHeaders = []string{"header\r\n"} },
		func(c *Config) { c.AllowedMethods = []string{"get"} },
		func(c *Config) { c.AllowedMethods = []string{"TRACE"} },
		func(c *Config) { c.MaxAge = -time.Second },
		func(c *Config) { c.MaxAge = time.Millisecond },
	} {
		c := exampleConfig()
		change(&c)
		if _, err := New(c); err == nil {
			t.Fatal(c)
		}
	}
	if err := (Config{}).Validate(); err != nil {
		t.Fatal(err)
	}
	for input, want := range map[string]string{"https://[::ffff:192.0.2.1]": "https://[::ffff:c000:201]", "https://[::ffff:c000:201]": "https://[::ffff:c000:201]", "https://ADMIN.EXAMPLE:443": "https://admin.example", "http://localhost:80": "http://localhost", "http://localhost:8080": "http://localhost:8080", "https://[::1]:443": "https://[::1]", "http://[::1]:8080": "http://[::1]:8080"} {
		got, ok := normalizeOrigin(input)
		if !ok || got != want {
			t.Fatalf("%q => %q %v", input, got, ok)
		}
	}
}

func TestPolicyCredentialsWildcardAndDuplicates(t *testing.T) {
	for _, credentials := range []bool{false, true} {
		c := exampleConfig()
		c.AllowCredentials = credentials
		if !credentials {
			c.AllowedOrigins = []string{"*"}
		}
		p, _ := New(c)
		for _, origin := range []string{"https://admin.example", "null", "https://admin.example/path", "https://bad\r\n"} {
			r := httptest.NewRequest("GET", "/", nil)
			r.Header.Set("Origin", origin)
			w := httptest.NewRecorder()
			p.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) })).ServeHTTP(w, r)
			if origin != "https://admin.example" {
				if w.Header().Get("Access-Control-Allow-Origin") != "" {
					t.Fatal(w.Header())
				}
				continue
			}
			want := "*"
			if credentials {
				want = origin
			}
			if w.Header().Get("Access-Control-Allow-Origin") != want || (w.Header().Get("Access-Control-Allow-Credentials") == "true") != credentials {
				t.Fatal(w.Header())
			}
		}
	}
	p, _ := New(exampleConfig())
	for _, header := range []string{"Origin", "Access-Control-Request-Method"} {
		r := httptest.NewRequest("OPTIONS", "/", nil)
		r.Header.Set("Origin", "https://admin.example")
		r.Header.Set("Access-Control-Request-Method", "POST")
		r.Header.Add(header, r.Header.Get(header))
		w := httptest.NewRecorder()
		p.Wrap(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { t.Fatal("ambiguous preflight reached handler") })).ServeHTTP(w, r)
		if w.Code != 403 || w.Header().Get("Access-Control-Allow-Origin") != "" {
			t.Fatal(w.Code, w.Header())
		}
	}
}

func TestPolicyPreservesHTTPAndSnapshot(t *testing.T) {
	c := exampleConfig()
	p, _ := New(c)
	c.AllowedOrigins[0] = "https://evil.example"
	c.AllowedHeaders[0] = "X-Evil"
	r := httptest.NewRequest("GET", "/", strings.NewReader("request"))
	r.Header.Set("Origin", "https://admin.example")
	type key struct{}
	r = r.WithContext(context.WithValue(r.Context(), key{}, "kept"))
	w := httptest.NewRecorder()
	w.Header().Set("Vary", "Accept-Encoding, origin")
	next := http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if writer != w || request != r || request.Context().Value(key{}) != "kept" {
			t.Fatal("HTTP capabilities replaced")
		}
		io.WriteString(writer, "chunk")
		writer.(http.Flusher).Flush()
	})
	p.Wrap(next).ServeHTTP(w, r)
	if !w.Flushed || w.Body.String() != "chunk" || w.Header().Get("Access-Control-Allow-Origin") != "https://admin.example" || len(w.Header().Values("Vary")) != 1 {
		t.Fatal(w)
	}
	for _, prior := range []string{"*", "Accept-Encoding, Origin"} {
		h := http.Header{"Vary": {prior}}
		vary(h, "Origin")
		if !reflect.DeepEqual(h.Values("Vary"), []string{prior}) {
			t.Fatal(h)
		}
	}
	disabled, _ := New(Config{})
	if reflect.ValueOf(disabled.Wrap(next)).Pointer() != reflect.ValueOf(next).Pointer() {
		t.Fatal("disabled policy not passthrough")
	}
	var wg sync.WaitGroup
	h := p.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) }))
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			for j := 0; j < 100; j++ {
				r := httptest.NewRequest("GET", "/", nil)
				origin := "https://admin.example"
				if i%2 == 1 {
					origin = "https://merchant.example"
				}
				r.Header.Set("Origin", origin)
				w := httptest.NewRecorder()
				h.ServeHTTP(w, r)
				if w.Header().Get("Access-Control-Allow-Origin") != origin {
					t.Error("cross-request origin leak")
				}
			}
		}(i)
	}
	wg.Wait()
}

func FuzzCORSHeaders(f *testing.F) {
	f.Add("https://admin.example", "POST", "authorization, content-type")
	f.Add("https://admin.example.evil", "POST", "")
	f.Add("null", "GET", "\r\n")
	p, _ := New(exampleConfig())
	h := p.Wrap(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic("preflight reached business") }))
	f.Fuzz(func(t *testing.T, origin, method, headers string) {
		r := httptest.NewRequest("OPTIONS", "/", nil)
		r.Header.Set("Origin", origin)
		r.Header.Set("Access-Control-Request-Method", method)
		if headers != "" {
			r.Header.Set("Access-Control-Request-Headers", headers)
		}
		w := httptest.NewRecorder()
		h.ServeHTTP(w, r)
		if w.Code != 204 && w.Code != 403 {
			t.Fatal(w.Code)
		}
		if w.Body.Len() != 0 {
			t.Fatal("unexpected body")
		}
		if value := w.Header().Get("Access-Control-Allow-Origin"); value != "" {
			if value != "https://admin.example" && value != "https://merchant.example" {
				t.Fatal("untrusted reflection")
			}
		}
	})
}

type benchmarkWriter struct{ header http.Header }

func (w benchmarkWriter) Header() http.Header         { return w.header }
func (w benchmarkWriter) Write(p []byte) (int, error) { return len(p), nil }
func (w benchmarkWriter) WriteHeader(int)             {}
func BenchmarkCORS(b *testing.B) {
	for _, mode := range []string{"disabled", "noOrigin", "allowed", "denied", "preflight", "wildcard"} {
		b.Run(mode, func(b *testing.B) {
			c := exampleConfig()
			if mode == "disabled" {
				c.Enabled = false
			}
			if mode == "wildcard" {
				c.AllowedOrigins = []string{"*"}
			}
			p, _ := New(c)
			h := p.Wrap(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) }))
			r := httptest.NewRequest("GET", "/", nil)
			r.Header.Set("Origin", "https://admin.example")
			if mode == "noOrigin" {
				r.Header.Del("Origin")
			}
			if mode == "denied" {
				r.Header.Set("Origin", "https://other.example")
			}
			if mode == "preflight" {
				r.Method = "OPTIONS"
				r.Header.Set("Access-Control-Request-Method", "POST")
				r.Header.Set("Access-Control-Request-Headers", "authorization, content-type")
			}
			w := benchmarkWriter{make(http.Header)}
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				clear(w.header)
				h.ServeHTTP(w, r)
			}
		})
	}
}