package main

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"
)

func opaqueReforgedTicket(fill byte) string {
	return "rft_" + strings.Repeat(string(fill), opaquePartBytes)
}

func TestValidReforgedTransportTicket(t *testing.T) {
	valid := opaqueReforgedTicket('a')
	if !validReforgedTransportTicket(valid) {
		t.Fatal("exact transport ticket was rejected")
	}
	for _, value := range []string{
		"", valid[:len(valid)-1], valid + "a",
		"rfi_" + valid[4:], "rfs_" + valid[4:],
		strings.Replace(valid, "a", ".", 1),
	} {
		if validReforgedTransportTicket(value) {
			t.Fatalf("unsafe transport ticket was accepted: %q", value)
		}
	}
}

func TestReforgedPublicRouteHasNoQuerySecret(t *testing.T) {
	if !validReforgedRouteQuery("") {
		t.Fatal("secret-free route was rejected")
	}
	for _, raw := range []string{
		"ticket=" + opaqueReforgedTicket('b'),
		"session=rfs_legacy", "x=1", "ticket=", "?",
	} {
		if validReforgedRouteQuery(raw) {
			t.Fatalf("public Reforged query was accepted: %q", raw)
		}
	}
}

func TestReforgedCredentialsAreRedactedAndClassicLabelIsUnchanged(t *testing.T) {
	lifecycle := "rfs_" + strings.Repeat("s", 43) +
		"rfa_" + strings.Repeat("a", 43)
	ticket := opaqueReforgedTicket('b')
	for _, secret := range []string{lifecycle, ticket} {
		if got := transportSessionLogLabel(secret, reforgedProfile); got != "[redacted]" || strings.Contains(got, secret) {
			t.Fatalf("Reforged credential entered log label: %q", got)
		}
	}
	classic := "classic-session-label"
	if got := transportSessionLogLabel(classic, ""); got != classic {
		t.Fatalf("Classic log label changed: %q", got)
	}
}

func TestReforgedPrivateRelayAuthenticatesInBandOnQuerylessPaths(t *testing.T) {
	ticket := opaqueReforgedTicket('q')
	auth, ok := reforgedRelayAuthMessage(ticket)
	if !ok {
		t.Fatal("valid relay ticket did not produce an auth frame")
	}
	expected := `{"type":"auth","profile":"` + reforgedProfile +
		`","ticket":"` + ticket + `"}`
	if string(auth) != expected {
		t.Fatalf("relay auth frame is not canonical: %q", auth)
	}
	for _, plane := range []string{"resource", "frame", "recovery", "input"} {
		target, valid := reforgedRelayTarget("ws://127.0.0.1:8245", plane)
		if !valid {
			t.Fatalf("valid relay plane was rejected: %s", plane)
		}
		parsed, err := url.Parse(target)
		if err != nil || parsed.RawQuery != "" || parsed.Fragment != "" ||
			strings.Contains(target, ticket) {
			t.Fatalf("relay target exposed a credential: %q err=%v", target, err)
		}
		ack := []byte(`{"type":"authOk","profile":"` + reforgedProfile +
			`","plane":"` + plane + `"}`)
		if !validReforgedRelayAuthAck(ack, plane) ||
			validReforgedRelayAuthAck(append(ack, '\n'), plane) {
			t.Fatalf("relay auth acknowledgement was not exact for %s", plane)
		}
	}
	if _, ok := reforgedRelayTarget("ws://127.0.0.1:8245", "signal"); ok {
		t.Fatal("public signaling plane entered the private bridge dialer")
	}
	if _, ok := reforgedRelayTarget("ws://relay.example:8245", "resource"); ok {
		t.Fatal("transport ticket could be sent to a non-loopback relay")
	}
	if _, ok := reforgedRelayAuthMessage("rfs_legacy"); ok {
		t.Fatal("lifecycle bearer produced a relay auth frame")
	}
}

func TestDialReforgedPlaneWaitsForExactAuthAcknowledgement(t *testing.T) {
	ticket := opaqueReforgedTicket('w')
	type observation struct {
		path  string
		query string
		kind  websocket.MessageType
		auth  string
		err   error
	}
	run := func(acknowledgement string) (observation, error) {
		seen := make(chan observation, 1)
		server := httptest.NewServer(http.HandlerFunc(func(
			writer http.ResponseWriter, request *http.Request,
		) {
			connection, err := websocket.Accept(writer, request, nil)
			if err != nil {
				seen <- observation{err: err}
				return
			}
			defer connection.CloseNow()
			kind, auth, err := connection.Read(request.Context())
			seen <- observation{path: request.URL.Path,
				query: request.URL.RawQuery, kind: kind, auth: string(auth), err: err}
			if err == nil {
				_ = connection.Write(request.Context(), websocket.MessageText,
					[]byte(acknowledgement))
			}
		}))
		defer server.Close()
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		plane, err := dialReforgedPlane(ctx,
			"ws"+strings.TrimPrefix(server.URL, "http"), "resource", ticket)
		if plane != nil {
			plane.conn.CloseNow()
		}
		return <-seen, err
	}

	authenticated := `{"type":"authOk","profile":"` + reforgedProfile +
		`","plane":"resource"}`
	seen, err := run(authenticated)
	if err != nil || seen.err != nil || seen.path != "/resource" ||
		seen.query != "" || seen.kind != websocket.MessageText {
		t.Fatalf("queryless relay authentication failed: seen=%+v err=%v",
			seen, err)
	}
	expectedAuth, _ := reforgedRelayAuthMessage(ticket)
	if seen.auth != string(expectedAuth) || strings.Contains(seen.path, ticket) {
		t.Fatalf("ticket was not confined to the first frame: seen=%+v", seen)
	}

	seen, err = run(`{"type":"authOk","profile":"` + reforgedProfile +
		`","plane":"frame"}`)
	if err == nil || seen.err != nil {
		t.Fatalf("wrong-plane relay acknowledgement was accepted: seen=%+v",
			seen)
	}
}

func TestReforgedHelloIsExactAndHasNoSpectatorMode(t *testing.T) {
	ticket := opaqueReforgedTicket('c')
	exact := `{"t":"hello","profile":"` + reforgedProfile +
		`","ticket":"` + ticket + `"}`
	if !validReforgedHello([]byte(exact), ticket, reforgedProfile) {
		t.Fatal("exact Reforged hello was rejected")
	}
	if parsed, ok := reforgedTransportTicketFromHello(
		[]byte(exact), reforgedProfile); !ok || parsed != ticket {
		t.Fatal("transport ticket was not extracted from exact hello")
	}
	for _, raw := range []string{
		`{"profile":"` + reforgedProfile + `","t":"hello","ticket":"` + ticket + `"}`,
		exact[:len(exact)-1] + `,"spectate":"token"}`,
		exact[:len(exact)-1] + `,"ticket":"` + ticket + `"}`,
		strings.Replace(exact, `"ticket"`, `"tick\u0065t"`, 1),
		exact + "\n",
	} {
		if validReforgedHello([]byte(raw), ticket, reforgedProfile) {
			t.Fatalf("non-canonical Reforged hello was accepted: %q", raw)
		}
		if _, ok := reforgedTransportTicketFromHello(
			[]byte(raw), reforgedProfile); ok {
			t.Fatalf("ticket was extracted from unsafe hello: %q", raw)
		}
	}
	if validReforgedHello([]byte(exact), opaqueReforgedTicket('d'),
		reforgedProfile) {
		t.Fatal("hello was accepted for another transport ticket")
	}
}

func TestValidateReforgedRouteConfig(t *testing.T) {
	if enabled, err := validateReforgedRouteConfig(0, 0); err != nil || enabled {
		t.Fatalf("zero/zero must leave Reforged disabled: enabled=%v err=%v",
			enabled, err)
	}
	if enabled, err := validateReforgedRouteConfig(8244, 9); err != nil || !enabled {
		t.Fatalf("valid isolated range was rejected: enabled=%v err=%v",
			enabled, err)
	}
	for _, values := range [][2]int{
		{8244, 0}, {0, 9}, {8244, 100}, {1023, 1}, {65535, 1},
	} {
		if enabled, err := validateReforgedRouteConfig(
			values[0], values[1]); err == nil || enabled {
			t.Fatalf("invalid config was accepted: base=%d slots=%d",
				values[0], values[1])
		}
	}
}

func TestReforgedRelayHostIsPinnedWithoutChangingClassicDefault(t *testing.T) {
	if err := validateReforgedRelayHost("127.0.0.1"); err != nil {
		t.Fatalf("exact loopback relay was rejected: %v", err)
	}
	for _, host := range []string{"localhost", "::1", "relay.example", ""} {
		if err := validateReforgedRelayHost(host); err == nil {
			t.Fatalf("unsafe Reforged relay host was accepted: %q", host)
		}
	}
}