package main

import (
	"encoding/json"
	"net"
	"net/http"
	"testing"
	"time"
)

// startHub spins up a hub with a given token and returns its address.
func startHub(t *testing.T, token string) (*peerRegistry, string) {
	t.Helper()
	reg := newPeerRegistry()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/ws", reg.wsHandler(token))
	srv := &http.Server{Handler: mux}
	go srv.Serve(ln)
	t.Cleanup(func() { srv.Close() })
	return reg, ln.Addr().String()
}

func peers(reg *peerRegistry) int { return len(reg.views()) }

// The correct token gets in and its data arrives.
func TestCorrectTokenGetsIn(t *testing.T) {
	reg, addr := startHub(t, "the-right-one")
	c, err := wsDial(addr, "/ws", 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if err := authenticate(c, config{Token: "the-right-one", clientID: "abc"}); err != nil {
		t.Fatalf("should have authenticated: %v", err)
	}
	msg, _ := json.Marshal(peerMessage{Type: "state", ID: "abc", Name: "Somebody", Deaths: 7})
	if err := c.WriteText(msg); err != nil {
		t.Fatal(err)
	}
	waitUntil(t, func() bool { return peers(reg) == 1 }, "the peer never showed up")
	if got := reg.views()[0].Deaths; got != 7 {
		t.Fatalf("expected 7 deaths, got %d", got)
	}
}

// The wrong token does NOT get in, and leaves no trace.
func TestWrongTokenDoesNotGetIn(t *testing.T) {
	reg, addr := startHub(t, "the-right-one")
	c, err := wsDial(addr, "/ws", 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	_ = authenticate(c, config{Token: "the-wrong-one", clientID: "x"})
	msg, _ := json.Marshal(peerMessage{Type: "state", ID: "x", Name: "Intruder", Deaths: 999})
	c.WriteText(msg)
	time.Sleep(700 * time.Millisecond)
	if peers(reg) != 0 {
		t.Fatalf("the hub accepted data with the wrong token: %+v", reg.views())
	}
}

// Sending data WITHOUT authenticating doesn't work either.
func TestUnauthenticatedDoesNotGetIn(t *testing.T) {
	reg, addr := startHub(t, "the-right-one")
	c, err := wsDial(addr, "/ws", 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	msg, _ := json.Marshal(peerMessage{Type: "state", ID: "x", Name: "Intruder", Deaths: 999})
	c.WriteText(msg)
	time.Sleep(700 * time.Millisecond)
	if peers(reg) != 0 {
		t.Fatalf("the hub accepted data with no authentication: %+v", reg.views())
	}
}

// The token never travels: the only thing that leaves the peer is an
// HMAC of the challenge.
func TestTheTokenNeverTravelsOverTheWire(t *testing.T) {
	_, addr := startHub(t, "very-secret-secret")
	c, err := wsDial(addr, "/ws", 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	raw, err := c.ReadMessage() // the challenge
	if err != nil {
		t.Fatal(err)
	}
	var ch authChallenge
	json.Unmarshal(raw, &ch)
	if ch.Nonce == "" {
		t.Fatal("the hub didn't send a challenge")
	}
	reply, _ := json.Marshal(authReply{Type: "auth", ID: "x", Proof: proofFor("very-secret-secret", ch.Nonce)})
	if string(reply) == "" {
		t.Fatal("empty reply")
	}
	for _, b := range [][]byte{raw, reply} {
		if contains(string(b), "very-secret-secret") {
			t.Fatalf("the token shows up in the traffic: %s", b)
		}
	}
}

// A captured response can't be replayed on another connection: the
// challenge changes.
func TestAnOldResponseCannotBeReplayed(t *testing.T) {
	reg, addr := startHub(t, "the-right-one")

	c1, err := wsDial(addr, "/ws", 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	raw, _ := c1.ReadMessage()
	var ch1 authChallenge
	json.Unmarshal(raw, &ch1)
	oldProof := proofFor("the-right-one", ch1.Nonce)
	c1.Close()

	c2, err := wsDial(addr, "/ws", 3*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c2.Close()
	raw2, _ := c2.ReadMessage()
	var ch2 authChallenge
	json.Unmarshal(raw2, &ch2)
	if ch2.Nonce == ch1.Nonce {
		t.Fatal("the challenge repeated across connections: the nonce has to differ")
	}
	reply, _ := json.Marshal(authReply{Type: "auth", ID: "x", Proof: oldProof})
	c2.WriteText(reply)
	msg, _ := json.Marshal(peerMessage{Type: "state", ID: "x", Name: "Replayer", Deaths: 1})
	c2.WriteText(msg)
	time.Sleep(700 * time.Millisecond)
	if peers(reg) != 0 {
		t.Fatalf("the hub accepted a replayed response: %+v", reg.views())
	}
}

// The generated token is random and of a reasonable length.
func TestGeneratedTokenIsStrong(t *testing.T) {
	testExeDir = t.TempDir()
	a, gen, err := resolveToken(config{})
	if err != nil || !gen {
		t.Fatalf("should have generated one: gen=%v err=%v", gen, err)
	}
	if len(a) != tokenBytes*2 {
		t.Fatalf("unexpected length: %d", len(a))
	}
	// the second time it reads it from the file, doesn't generate another
	b, gen2, _ := resolveToken(config{})
	if gen2 || a != b {
		t.Fatalf("should reuse the saved one: %q vs %q (gen=%v)", a, b, gen2)
	}
	// and a different installation gets a different one
	testExeDir = t.TempDir()
	c, _, _ := resolveToken(config{})
	if c == a {
		t.Fatal("two installations generated the same token")
	}
	testExeDir = ""
}

func contains(s, sub string) bool {
	return len(sub) > 0 && len(s) >= len(sub) && (func() bool {
		for i := 0; i+len(sub) <= len(s); i++ {
			if s[i:i+len(sub)] == sub {
				return true
			}
		}
		return false
	})()
}

func waitUntil(t *testing.T, cond func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(4 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal(msg)
}