package host

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"testing/fstest"
)

// upstream records what reached the homeserver.
type upstream struct {
	srv  *httptest.Server
	hits []*http.Request
}

func newUpstream(t *testing.T) *upstream {
	t.Helper()
	u := &upstream{}
	u.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u.hits = append(u.hits, r.Clone(r.Context()))
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Set-Cookie", "tracker=1")
		if strings.HasPrefix(r.URL.Path, "/_matrix/client/v1/media/") {
			w.Header().Set("Cache-Control", "private, max-age=86400")
		}
		_, _ = w.Write([]byte(`{"path":"` + r.URL.EscapedPath() + `"}`))
	}))
	t.Cleanup(u.srv.Close)
	return u
}

var clientFS = fstest.MapFS{
	"index.html":         {Data: []byte("<!doctype html><title>client</title>")},
	"favicon.svg":        {Data: []byte("<svg/>")},
	"assets/index-a1.js": {Data: []byte("console.log(1)")},
}

func newHandler(t *testing.T, up *upstream, built bool) http.Handler {
	t.Helper()
	target, _ := url.Parse(up.srv.URL)
	opts := Options{Homeserver: target, Version: "1.2.3"}
	if built {
		opts.Client = clientFS
	}
	h, err := New(opts)
	if err != nil {
		t.Fatal(err)
	}
	return h
}

func do(h http.Handler, method, target string, header http.Header) *httptest.ResponseRecorder {
	req := httptest.NewRequest(method, target, nil)
	for k, v := range header {
		req.Header[k] = v
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestProxiesClientAndMediaPaths(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	cases := []struct{ method, path string }{
		{http.MethodGet, "/_matrix/client/versions"},
		{http.MethodPost, "/_matrix/client/v3/login"},
		{http.MethodGet, "/_matrix/client/v3/sync?timeout=30000&since=s1_2"},
		{http.MethodPut, "/_matrix/client/v3/rooms/%21a%3Awarp.example/send/m.room.message/txn1"},
		{http.MethodPost, "/_matrix/media/v3/upload?filename=cat.png"},
		{http.MethodGet, "/_matrix/client/v1/media/download/warp.example/abc"},
		{http.MethodGet, "/_matrix/client/v3/pushrules/"},
	}
	for _, c := range cases {
		rec := do(h, c.method, c.path, nil)
		if rec.Code != http.StatusOK {
			t.Errorf("%s %s: status %d, want 200", c.method, c.path, rec.Code)
		}
	}
	if len(up.hits) != len(cases) {
		t.Fatalf("upstream saw %d requests, want %d", len(up.hits), len(cases))
	}
	// Escaped IDs and the query reach the homeserver unchanged.
	if got := up.hits[3].URL.EscapedPath(); got != "/_matrix/client/v3/rooms/%21a%3Awarp.example/send/m.room.message/txn1" {
		t.Errorf("escaped path = %q", got)
	}
	if got := up.hits[2].URL.RawQuery; got != "timeout=30000&since=s1_2" {
		t.Errorf("query = %q", got)
	}
}

func TestRefusesOtherHomeserverPaths(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	for _, p := range []string{
		"/_matrix/federation/v1/version",
		"/_matrix/key/v2/server",
		"/_matrix/app/v1/transactions/1",
		"/_warp/admin/v1/stats",
		"/_synapse/admin/v1/users",
		"/_matrix/client/../federation/v1/version",
		"/_matrix/client/v3/../../federation/v1/version",
		"/_matrix/client//v3/sync",
		"/_matrix/client/v3/pushrules//",
		"/_matrix/client/v3/../../federation/",
	} {
		rec := do(h, http.MethodGet, p, nil)
		if rec.Code != http.StatusNotFound {
			t.Errorf("%s: status %d, want 404", p, rec.Code)
		}
		var body map[string]string
		_ = json.Unmarshal(rec.Body.Bytes(), &body)
		if body["errcode"] != "M_UNRECOGNIZED" {
			t.Errorf("%s: errcode %q", p, body["errcode"])
		}
	}
	if len(up.hits) != 0 {
		t.Fatalf("upstream saw %d requests, want none", len(up.hits))
	}
}

func TestProxyStripsCookiesAndForwardsToken(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	rec := do(h, http.MethodGet, "/_matrix/client/v3/account/whoami", http.Header{
		"Authorization": {"Bearer tok"},
		"Cookie":        {"session=secret"},
	})
	if rec.Header().Get("Set-Cookie") != "" {
		t.Error("Set-Cookie from the homeserver was passed to the browser")
	}
	if rec.Header().Get("Cache-Control") != "no-store" {
		t.Errorf("Cache-Control = %q, want no-store", rec.Header().Get("Cache-Control"))
	}
	got := up.hits[0]
	if got.Header.Get("Cookie") != "" {
		t.Error("Cookie was passed to the homeserver")
	}
	if got.Header.Get("Authorization") != "Bearer tok" {
		t.Errorf("Authorization = %q", got.Header.Get("Authorization"))
	}
	if got.Header.Get("X-Forwarded-For") == "" {
		t.Error("X-Forwarded-For not set")
	}

	// The homeserver's own caching of media is kept.
	rec = do(h, http.MethodGet, "/_matrix/client/v1/media/download/warp.example/abc", nil)
	if got := rec.Header().Get("Cache-Control"); got != "private, max-age=86400" {
		t.Errorf("media Cache-Control = %q", got)
	}
}

func TestUnreachableHomeserverIs502(t *testing.T) {
	target, _ := url.Parse("http://127.0.0.1:1")
	h, err := New(Options{Homeserver: target, Client: clientFS})
	if err != nil {
		t.Fatal(err)
	}
	rec := do(h, http.MethodGet, "/_matrix/client/versions", nil)
	if rec.Code != http.StatusBadGateway {
		t.Fatalf("status %d, want 502", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "M_UNKNOWN") {
		t.Errorf("body %q", rec.Body.String())
	}
}

func TestServesClient(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	cases := []struct {
		path, wantBody, wantCache string
		wantCode                  int
	}{
		{"/", "client", "no-cache", 200},
		{"/room/!abc:warp.example", "client", "no-cache", 200},
		{"/room/%21abc%3Awarp.example", "client", "no-cache", 200},
		{"/index.html", "client", "no-cache", 200},
		{"/assets/index-a1.js", "console.log", "public, max-age=31536000, immutable", 200},
		{"/favicon.svg", "<svg", "public, max-age=3600", 200},
		{"/assets/gone-b2.js", "", "", 404},
		{"/robots.txt", "", "", 404},
	}
	for _, c := range cases {
		rec := do(h, http.MethodGet, c.path, nil)
		if rec.Code != c.wantCode {
			t.Errorf("%s: status %d, want %d", c.path, rec.Code, c.wantCode)
			continue
		}
		if c.wantCode != 200 {
			continue
		}
		body, _ := io.ReadAll(rec.Body)
		if !strings.Contains(string(body), c.wantBody) {
			t.Errorf("%s: body %q", c.path, body)
		}
		if got := rec.Header().Get("Cache-Control"); got != c.wantCache {
			t.Errorf("%s: Cache-Control %q, want %q", c.path, got, c.wantCache)
		}
		if csp := rec.Header().Get("Content-Security-Policy"); !strings.Contains(csp, "connect-src 'self';") {
			t.Errorf("%s: CSP %q", c.path, csp)
		}
		// WebAssembly for end-to-end encryption, and no JavaScript eval.
		if csp := rec.Header().Get("Content-Security-Policy"); !strings.Contains(csp, "script-src 'self' 'wasm-unsafe-eval';") {
			t.Errorf("%s: CSP %q lacks wasm-unsafe-eval", c.path, csp)
		}
		if rec.Header().Get("X-Frame-Options") != "DENY" {
			t.Errorf("%s: X-Frame-Options %q", c.path, rec.Header().Get("X-Frame-Options"))
		}
	}

	if rec := do(h, http.MethodPost, "/", nil); rec.Code != http.StatusMethodNotAllowed {
		t.Errorf("POST /: status %d, want 405", rec.Code)
	}
}

func TestConfigProxied(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)

	rec := do(h, http.MethodGet, "/config.json", nil)
	if rec.Code != 200 {
		t.Fatalf("status %d", rec.Code)
	}
	var cfg Config
	if err := json.Unmarshal(rec.Body.Bytes(), &cfg); err != nil {
		t.Fatal(err)
	}
	if cfg.HomeserverURL != "" || cfg.Version != "1.2.3" {
		t.Errorf("config = %+v, want same-origin and the version", cfg)
	}
	if rec.Header().Get("Cache-Control") != "no-cache" {
		t.Errorf("Cache-Control %q", rec.Header().Get("Cache-Control"))
	}
}

func TestDirectModeProxiesNothing(t *testing.T) {
	public, _ := url.Parse("https://matrix.warp.example")
	h, err := New(Options{PublicHomeserver: public, Client: clientFS})
	if err != nil {
		t.Fatal(err)
	}

	rec := do(h, http.MethodGet, "/config.json", nil)
	var cfg Config
	_ = json.Unmarshal(rec.Body.Bytes(), &cfg)
	if cfg.HomeserverURL != "https://matrix.warp.example" {
		t.Errorf("homeserver_url = %q", cfg.HomeserverURL)
	}

	if rec := do(h, http.MethodGet, "/_matrix/client/versions", nil); rec.Code != http.StatusNotFound {
		t.Errorf("proxy in direct mode: status %d, want 404", rec.Code)
	}

	// The browser may now connect to the homeserver, and only there.
	rec = do(h, http.MethodGet, "/", nil)
	if csp := rec.Header().Get("Content-Security-Policy"); !strings.Contains(csp, "connect-src 'self' https://matrix.warp.example;") {
		t.Errorf("CSP %q", csp)
	}
}

func TestNotBuiltClient(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, false)

	rec := do(h, http.MethodGet, "/", nil)
	if rec.Code != http.StatusServiceUnavailable || !strings.Contains(rec.Body.String(), "isn’t built") {
		t.Fatalf("status %d body %q", rec.Code, rec.Body.String())
	}
	// The proxy still works without the client.
	if rec := do(h, http.MethodGet, "/_matrix/client/versions", nil); rec.Code != http.StatusOK {
		t.Errorf("proxy status %d", rec.Code)
	}
}

func TestHealthz(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)
	if rec := do(h, http.MethodGet, "/healthz", nil); rec.Code != 200 || rec.Body.String() != "ok\n" {
		t.Fatalf("status %d body %q", rec.Code, rec.Body.String())
	}
}

func TestNewValidatesHomeserver(t *testing.T) {
	rel, _ := url.Parse("warp:8008")
	ftp, _ := url.Parse("ftp://warp.example")
	ok, _ := url.Parse("http://warp:8008")
	for name, opts := range map[string]Options{
		"relative": {Homeserver: rel},
		"scheme":   {PublicHomeserver: ftp},
		"neither":  {},
		"both":     {Homeserver: ok, PublicHomeserver: ok},
	} {
		if _, err := New(opts); err == nil {
			t.Errorf("%s: want an error", name)
		}
	}
}

func TestProxiesWellKnownClient(t *testing.T) {
	up := newUpstream(t)
	h := newHandler(t, up, true)
	rec := do(h, http.MethodGet, "/.well-known/matrix/client", nil)
	if rec.Code != 200 || len(up.hits) != 1 || up.hits[0].URL.Path != "/.well-known/matrix/client" {
		t.Fatalf("status %d, hits %d", rec.Code, len(up.hits))
	}
	// The server delegation document stays private.
	if rec := do(h, http.MethodGet, "/.well-known/matrix/server", nil); rec.Code == 200 && len(up.hits) > 1 {
		t.Errorf("/.well-known/matrix/server was proxied")
	}
}

func TestCallOriginsJoinConnectSrc(t *testing.T) {
	up := newUpstream(t)
	target, _ := url.Parse(up.srv.URL)
	h, err := New(Options{Homeserver: target, Client: clientFS,
		CallOrigins: []string{"https://livekit-jwt.warp.example/", "wss://livekit.warp.example"}})
	if err != nil {
		t.Fatal(err)
	}
	rec := do(h, http.MethodGet, "/", nil)
	want := "connect-src 'self' https://livekit-jwt.warp.example wss://livekit.warp.example https://livekit.warp.example;"
	if csp := rec.Header().Get("Content-Security-Policy"); !strings.Contains(csp, want) {
		t.Errorf("CSP %q", csp)
	}

	if _, err := New(Options{Homeserver: target, CallOrigins: []string{"https://x.example/path"}}); err == nil {
		t.Error("a call origin with a path was accepted")
	}
}