SimpleRemoter/server/go/wsauth/wsauth_test.go

package wsauth

import (
	"testing"
	"time"
)

func TestSHA256Vector(t *testing.T) {
	// Known vector — keeps us honest against accidental algorithm changes.
	got := ComputeSHA256("abc")
	want := "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
	if got != want {
		t.Fatalf("SHA256(abc): got %s want %s", got, want)
	}
}

// adminLoginResponse helps tests compute the right login response for an
// admin account that now uses a real per-instance salt.
func adminLoginResponse(t *testing.T, a *Authenticator, username, password, nonce string) string {
	t.Helper()
	salt, ok := a.GetSalt(username)
	if !ok {
		t.Fatalf("admin %s not registered", username)
	}
	return ComputeSHA256(HashPassword(password, salt) + nonce)
}

func TestLoginRoundTripAdmin(t *testing.T) {
	a := New()
	a.AddAdminFromPlainPassword("admin", "hunter2")

	salt, ok := a.GetSalt("admin")
	if !ok {
		t.Fatal("admin should be found")
	}
	if len(salt) != 2*saltBytes {
		t.Fatalf("admin salt should be a real 16-hex value, got %q (len=%d)", salt, len(salt))
	}

	nonce := "abc123"
	response := adminLoginResponse(t, a, "admin", "hunter2", nonce)

	token, role, err := a.VerifyLogin("admin", response, nonce)
	if err != nil {
		t.Fatalf("VerifyLogin: %v", err)
	}
	if role != "admin" {
		t.Fatalf("role: got %q want admin", role)
	}
	if len(token) != 2*tokenBytes {
		t.Fatalf("token length: got %d want %d", len(token), 2*tokenBytes)
	}

	sess, err := a.ValidateToken(token)
	if err != nil {
		t.Fatalf("ValidateToken: %v", err)
	}
	if sess.Username != "admin" || sess.Role != "admin" {
		t.Fatalf("session: %+v", sess)
	}
}

func TestLoginRoundTripViewerWithSalt(t *testing.T) {
	a := New()
	salt, _ := NewSalt()
	a.AddUser(User{
		Username:     "alice",
		PasswordHash: HashPassword("p@ss", salt),
		Salt:         salt,
		Role:         "viewer",
	})

	gotSalt, ok := a.GetSalt("alice")
	if !ok || gotSalt != salt {
		t.Fatalf("salt: ok=%v got=%q want=%q", ok, gotSalt, salt)
	}

	nonce, _ := NewNonce()
	response := ComputeSHA256(HashPassword("p@ss", salt) + nonce)
	_, role, err := a.VerifyLogin("alice", response, nonce)
	if err != nil || role != "viewer" {
		t.Fatalf("VerifyLogin: role=%q err=%v", role, err)
	}
}

// TestGetSaltUnknownUserShape verifies the salt-probe mitigation: an
// unknown user must get back a value that's shape-identical to a real
// salt, so an attacker can't tell from /get_salt alone whether a
// username exists.
func TestGetSaltUnknownUserShape(t *testing.T) {
	a := New()
	a.AddAdminFromPlainPassword("admin", "pw")

	fake, ok := a.GetSalt("nobody")
	if ok {
		t.Fatal("ok should be false for unknown user")
	}
	if len(fake) != 2*saltBytes {
		t.Fatalf("fake salt should be %d hex chars; got %q (len=%d)", 2*saltBytes, fake, len(fake))
	}
	// Determinism: repeated probes for the same username get the same fake.
	fake2, _ := a.GetSalt("nobody")
	if fake != fake2 {
		t.Fatalf("fake salt should be deterministic for repeated probes; got %q vs %q", fake, fake2)
	}
	// Different usernames get different fake salts.
	other, _ := a.GetSalt("ghost")
	if fake == other {
		t.Fatalf("fake salts should differ across usernames; both = %q", fake)
	}
}

func TestLoginRejectsWrongResponse(t *testing.T) {
	a := New()
	a.AddAdminFromPlainPassword("admin", "x")
	_, _, err := a.VerifyLogin("admin", "deadbeef", "nonce")
	if err == nil {
		t.Fatal("expected error for bad response")
	}
	_, _, err = a.VerifyLogin("ghost", "anything", "anything")
	if err == nil {
		t.Fatal("expected error for unknown user")
	}
}

func TestTokenExpiry(t *testing.T) {
	a := New()
	a.SetTokenExpire(50 * time.Millisecond)
	a.AddAdminFromPlainPassword("admin", "x")
	nonce, _ := NewNonce()
	response := adminLoginResponse(t, a, "admin", "x", nonce)
	token, _, err := a.VerifyLogin("admin", response, nonce)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := a.ValidateToken(token); err != nil {
		t.Fatalf("fresh token should validate: %v", err)
	}
	time.Sleep(80 * time.Millisecond)
	if _, err := a.ValidateToken(token); err == nil {
		t.Fatal("expired token should not validate")
	}
}

func TestRevoke(t *testing.T) {
	a := New()
	a.AddAdminFromPlainPassword("admin", "x")
	nonce, _ := NewNonce()
	response := adminLoginResponse(t, a, "admin", "x", nonce)
	token, _, _ := a.VerifyLogin("admin", response, nonce)
	a.RevokeToken(token)
	if _, err := a.ValidateToken(token); err == nil {
		t.Fatal("revoked token should not validate")
	}
}

func TestRateLimiterAllowsBurstThenBlocks(t *testing.T) {
	r := NewRateLimiter(3, time.Minute)
	for i := 0; i < 3; i++ {
		if !r.Allow("ip-a") {
			t.Fatalf("attempt %d should be allowed", i+1)
		}
	}
	if r.Allow("ip-a") {
		t.Fatal("4th attempt should be denied")
	}
	// Different key has independent budget.
	if !r.Allow("ip-b") {
		t.Fatal("different key should still be allowed")
	}
}

func TestRateLimiterReset(t *testing.T) {
	r := NewRateLimiter(2, time.Minute)
	r.Allow("k")
	r.Allow("k")
	if r.Allow("k") {
		t.Fatal("3rd should be denied")
	}
	r.Reset("k")
	if !r.Allow("k") {
		t.Fatal("after Reset, should be allowed again")
	}
}

func TestRateLimiterDisabledWhenZeroLimit(t *testing.T) {
	r := NewRateLimiter(0, time.Minute)
	for i := 0; i < 100; i++ {
		if !r.Allow("k") {
			t.Fatalf("limit=0 should never deny, denied at i=%d", i)
		}
	}
}

func TestRateLimiterNilSafe(t *testing.T) {
	var r *RateLimiter
	if !r.Allow("anything") {
		t.Fatal("nil limiter should allow")
	}
	r.Reset("anything")
	r.Sweep()
}