package main

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"
)

// fakeProvider is discovery, a key set, and a token endpoint. A real provider
// is checked by hand.
type fakeProvider struct {
	*httptest.Server
	key      *rsa.PrivateKey
	signer   *rsa.PrivateKey // what it signs with, which is key unless a test says otherwise
	claims   map[string]any
	verifier string // the code_verifier the token endpoint was given
}

func newFakeProvider(t *testing.T) *fakeProvider {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	p := &fakeProvider{key: key, signer: key}
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]string{
			"issuer":                 p.URL,
			"authorization_endpoint": p.URL + "/auth",
			"token_endpoint":         p.URL + "/token",
			"jwks_uri":               p.URL + "/jwks",
		})
	})
	mux.HandleFunc("/jwks", func(w http.ResponseWriter, r *http.Request) {
		pub := key.Public().(*rsa.PublicKey)
		json.NewEncoder(w).Encode(map[string]any{"keys": []map[string]string{{
			"kty": "RSA", "kid": "k1", "alg": "RS256", "use": "sig",
			"n": base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
			"e": base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes()),
		}}})
	})
	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		p.verifier = r.PostFormValue("code_verifier")
		json.NewEncoder(w).Encode(map[string]string{"id_token": p.sign(t)})
	})
	p.Server = httptest.NewServer(mux)
	p.claims = map[string]any{
		"iss": p.URL, "sub": "3719", "aud": "git-forge",
		"exp": time.Now().Add(time.Hour).Unix(), "iat": time.Now().Unix(),
		"email": "carol@example.com", "email_verified": true,
		"preferred_username": "Carol",
	}
	return p
}

func (p *fakeProvider) sign(t *testing.T) string {
	t.Helper()
	b64 := base64.RawURLEncoding.EncodeToString
	header, err := json.Marshal(map[string]string{"alg": "RS256", "kid": "k1", "typ": "JWT"})
	if err != nil {
		t.Fatal(err)
	}
	payload, err := json.Marshal(p.claims)
	if err != nil {
		t.Fatal(err)
	}
	signed := b64(header) + "." + b64(payload)
	sum := sha256.Sum256([]byte(signed))
	sig, err := rsa.SignPKCS1v15(rand.Reader, p.signer, crypto.SHA256, sum[:])
	if err != nil {
		t.Fatal(err)
	}
	return signed + "." + b64(sig)
}

// login drives one whole flow and returns the callback's response.
func login(t *testing.T, f *Forge, h http.Handler, p *fakeProvider, mangle func(url.Values)) *httptest.ResponseRecorder {
	t.Helper()
	start := get(h, "/oidc/sso", nil)
	if start.Code != http.StatusSeeOther {
		t.Fatalf("starting a login: status %d", start.Code)
	}
	to, err := url.Parse(start.Header().Get("Location"))
	if err != nil {
		t.Fatal(err)
	}
	sent := to.Query()
	if sent.Get("code_challenge_method") != "S256" || sent.Get("code_challenge") == "" {
		t.Errorf("no pkce challenge in %v", sent)
	}

	back := url.Values{"code": {"the-code"}, "state": {sent.Get("state")}}
	if mangle != nil {
		mangle(back)
	}
	w := get(h, "/oidc/sso/callback?"+back.Encode(), start.Result().Cookies())

	// The verifier the token endpoint saw has to hash to the challenge sent.
	if p.verifier != "" {
		sum := sha256.Sum256([]byte(p.verifier))
		if got := base64.RawURLEncoding.EncodeToString(sum[:]); got != sent.Get("code_challenge") {
			t.Errorf("the verifier does not match the challenge")
		}
	}
	return w
}

func oidcForge(t *testing.T, p *fakeProvider) *Forge {
	t.Helper()
	f := testForge(t)
	// Assigned rather than loaded: the config refuses an issuer that is not
	// https, and this one is a test server.
	f.cfg.OIDC = []OIDCProvider{{
		Key: "sso", Label: "Example SSO", Issuer: p.URL,
		ClientID: "git-forge", ClientSecret: "shh",
	}}
	return f
}

func TestOIDCLoginProvisions(t *testing.T) {
	p := newFakeProvider(t)
	defer p.Close()
	f := oidcForge(t, p)
	h := f.routes()

	w := login(t, f, h, p, nil)
	if w.Code != http.StatusSeeOther {
		t.Fatalf("callback: status %d, %s", w.Code, w.Body.String())
	}
	u, ok := f.store.User("sso:carol")
	if !ok {
		t.Fatalf("no account was provisioned: %v", f.store.Users())
	}
	if u.Origin != originOIDC || u.Iss != p.URL || u.Sub != "3719" {
		t.Errorf("identity is %+v", u)
	}
	if u.Email != "carol@example.com" || !u.EmailConfirmed {
		t.Errorf("address is %q, confirmed %v", u.Email, u.EmailConfirmed)
	}
	if u.Admin || u.Password != "" {
		t.Error("a provisioned account should hold no grants and no password")
	}
	if body := get(h, "/", w.Result().Cookies()).Body.String(); !strings.Contains(body, "sso:carol") {
		t.Error("the session from an external login does not work")
	}

	// A second login is the same account: identity is iss plus sub, and the
	// name is never matched on.
	p.claims["preferred_username"] = "Carol Renamed"
	if w := login(t, f, h, p, nil); w.Code != http.StatusSeeOther {
		t.Fatalf("second login: status %d", w.Code)
	}
	if len(f.store.Users()) != 1 {
		t.Errorf("a second login made a second account: %v", f.store.Users())
	}

	// A different subject with the same username gets a suffix rather than
	// the other account.
	p.claims["sub"] = "9999"
	p.claims["preferred_username"] = "Carol"
	p.claims["email"] = "carol2@example.com"
	if w := login(t, f, h, p, nil); w.Code != http.StatusSeeOther {
		t.Fatalf("third login: status %d", w.Code)
	}
	if _, ok := f.store.User("sso:carol2"); !ok {
		t.Errorf("a colliding name was not suffixed: %v", f.store.Users())
	}
}

// The provider is the first factor and not the second, and the global admin
// flag rests on there being a second one.
func TestOIDCLoginAsksForTheSecondFactor(t *testing.T) {
	p := newFakeProvider(t)
	defer p.Close()
	f := oidcForge(t, p)
	h := f.routes()

	if w := login(t, f, h, p, nil); w.Code != http.StatusSeeOther {
		t.Fatalf("the login that provisions the account: status %d", w.Code)
	}
	addTOTP(t, f, "sso:carol")

	first := login(t, f, h, p, nil)
	if first.Code != http.StatusOK || !strings.Contains(first.Body.String(), "Authenticator code") {
		t.Fatalf("the callback did not ask for a code: status %d", first.Code)
	}
	for _, c := range first.Result().Cookies() {
		if c.Name == sessionCookie && c.Value != "" {
			t.Fatal("the provider alone issued a session")
		}
	}

	second := post(h, "/login/totp", url.Values{
		"csrf":    {f.csrf(withCookies("/", nil))},
		"name":    {between(first.Body.String(), `name="name" value="`, `"`)},
		"expires": {between(first.Body.String(), `name="expires" value="`, `"`)},
		"t":       {between(first.Body.String(), `name="t" value="`, `"`)},
		"code":    {totpCode(knownSecret, time.Now())},
	}, nil)
	if second.Code != http.StatusSeeOther {
		t.Fatalf("the second step refused a good code: status %d", second.Code)
	}
	if body := get(h, "/", second.Result().Cookies()).Body.String(); !strings.Contains(body, "sso:carol") {
		t.Error("the session the second step issued does not work")
	}
}

// The note carrying a half-finished login is bound to the epoch, so signing the
// account out everywhere while it is in flight kills it.
func TestPendingLoginDiesWithTheEpoch(t *testing.T) {
	f := testForge(t)
	addUser(t, f, "alice", "hunter2hunter2")
	addTOTP(t, f, "alice")
	h := f.routes()
	first := post(h, "/login", url.Values{"name": {"alice"}, "password": {"hunter2hunter2"}}, nil)
	pending := url.Values{
		"csrf":    {f.csrf(withCookies("/", nil))},
		"name":    {between(first.Body.String(), `name="name" value="`, `"`)},
		"expires": {between(first.Body.String(), `name="expires" value="`, `"`)},
		"t":       {between(first.Body.String(), `name="t" value="`, `"`)},
		"code":    {totpCode(knownSecret, time.Now())},
	}
	err := f.store.UpdateUsers("users: sign out", f.store.forgeAuthor(), func(m map[string]User) error {
		u := m["alice"]
		u.Epoch++
		m["alice"] = u
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if w := post(h, "/login/totp", pending, nil); w.Code != http.StatusForbidden {
		t.Errorf("a login in flight survived the account being signed out: status %d", w.Code)
	}
}

func TestOIDCRefusals(t *testing.T) {
	cases := []struct {
		what   string
		claims map[string]any
		mangle func(url.Values)
	}{
		{what: "an unverified address", claims: map[string]any{"email_verified": false}},
		{what: "a missing email_verified", claims: map[string]any{"email_verified": nil}},
		{what: "an id token for somebody else", claims: map[string]any{"aud": "another-client"}},
		{what: "an expired id token", claims: map[string]any{"exp": time.Now().Add(-2 * time.Hour).Unix()}},
		{what: "an id token from another issuer", claims: map[string]any{"iss": "https://evil.example"}},
		{what: "an address with no domain", claims: map[string]any{"email": "carol"}},
		{what: "an address with a line break in it", claims: map[string]any{"email": "carol@example.com\nx@evil.example"}},
		{what: "a state that did not come from us", mangle: func(v url.Values) { v.Set("state", "made-up") }},
	}
	for _, c := range cases {
		p := newFakeProvider(t)
		f := oidcForge(t, p)
		for k, v := range c.claims {
			if v == nil {
				delete(p.claims, k)
			} else {
				p.claims[k] = v
			}
		}
		w := login(t, f, f.routes(), p, c.mangle)
		if w.Code != http.StatusForbidden {
			t.Errorf("%s: status %d", c.what, w.Code)
		}
		if len(f.store.Users()) != 0 {
			t.Errorf("%s: an account was provisioned anyway", c.what)
		}
		p.Close()
	}
}

// Every claim above is read out of a payload that has already been trusted, so
// a token the provider's own key set cannot vouch for has to be refused before
// any of them is looked at. The key set still publishes k1 and the token still
// names it: only the signature is somebody else's.
func TestAnIdTokenTheKeySetCannotVouchForIsRefused(t *testing.T) {
	p := newFakeProvider(t)
	defer p.Close()
	other, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	p.signer = other

	f := oidcForge(t, p)
	if w := login(t, f, f.routes(), p, nil); w.Code != http.StatusForbidden {
		t.Errorf("a token signed by a key the provider does not publish: status %d", w.Code)
	}
	if len(f.store.Users()) != 0 {
		t.Errorf("an account was provisioned from a token nothing vouches for: %v", f.store.Users())
	}
}

// An exact kid wins: taking a key that merely names none ahead of the match
// fails verification and hides a rotation from the refetch.
func TestFindKey(t *testing.T) {
	mixed := []jwk{{Kid: ""}, {Kid: "k2"}}
	if k, ok := findKey(mixed, "k2"); !ok || k.Kid != "k2" {
		t.Errorf("an exact kid was passed over for %+v, %v", k, ok)
	}
	if k, ok := findKey(mixed, ""); !ok || k.Kid != "" {
		t.Errorf("a token naming no kid got %+v, %v", k, ok)
	}
	if k, ok := findKey([]jwk{{Kid: "sole", Use: "sig"}}, ""); !ok || k.Kid != "sole" {
		t.Errorf("the one signing key was refused: %+v, %v", k, ok)
	}
	// A rotation: every key names a kid and none of them is this one, so the
	// answer has to be no, or the key set is never fetched again.
	if _, ok := findKey([]jwk{{Kid: "k1"}, {Kid: "k2"}}, "k9"); ok {
		t.Error("an unknown kid found a key anyway")
	}
	if _, ok := findKey([]jwk{{Kid: "k1", Use: "enc"}}, ""); ok {
		t.Error("an encryption key was offered for a signature")
	}
}

func TestSanitizeName(t *testing.T) {
	cases := map[string]string{
		"Carol":                 "carol",
		"carol.smith":           "carol.smith",
		"../../config":          "config",
		"Ünïcodé":               "n-cod", // no transliteration: unsafe runs become a dash
		"":                      "",
		"...":                   "",
		"a b":                   "a-b",
		"sso:carol":             "sso-carol",
		strings.Repeat("x", 40): "",
	}
	for in, want := range cases {
		if got := sanitizeName(in); got != want {
			t.Errorf("sanitizeName(%q) = %q, wanted %q", in, got, want)
		}
	}
}

func TestOIDCAllows(t *testing.T) {
	cases := []struct {
		allow []string
		email string
		want  bool
	}{
		{nil, "anyone@example.com", true},
		{[]string{}, "anyone@example.com", true},
		{[]string{"carol@example.com"}, "carol@example.com", true},
		{[]string{"carol@example.com"}, "dave@example.com", false},
		{[]string{"@example.com"}, "anyone@example.com", true},
		{[]string{"@example.com"}, "anyone@evil.example", false},
		// A domain entry is a whole label, not a suffix.
		{[]string{"@example.com"}, "anyone@sub.example.com", false},
		{[]string{"@example.com"}, "anyone@example.com.attacker.test", false},
		{[]string{"@example.com"}, "anyone@notexample.com", false},
		// The provider decides the case of what it sends.
		{[]string{"Carol@Example.COM"}, "carol@example.com", true},
		{[]string{"@Example.COM"}, "CAROL@example.com", true},
		{[]string{"@example.com", "dave@other.example"}, "dave@other.example", true},
		{[]string{"@example.com"}, "no-at-sign-at-all", false},
	}
	for _, c := range cases {
		p := OIDCProvider{Allow: c.allow}
		if got := p.allows(c.email); got != c.want {
			t.Errorf("allow %v, address %q: got %v", c.allow, c.email, got)
		}
	}
}

// An address outside the list is refused before an account is made.
func TestOIDCLoginRefusesUnlistedAddress(t *testing.T) {
	p := newFakeProvider(t)
	defer p.Close()
	f := oidcForge(t, p)
	f.cfg.OIDC[0].Allow = []string{"@allowed.example", "specific@other.example"}

	if w := login(t, f, f.routes(), p, nil); w.Code != http.StatusForbidden {
		t.Fatalf("an unlisted address got status %d", w.Code)
	}
	if n := len(f.store.Users()); n != 0 {
		t.Fatalf("%d accounts were provisioned anyway: %v", n, f.store.Users())
	}

	// And the listed one gets in, so the refusal is the list and not the flow.
	p.claims["email"] = "carol@allowed.example"
	if w := login(t, f, f.routes(), p, nil); w.Code != http.StatusSeeOther {
		t.Fatalf("a listed address got status %d", w.Code)
	}
	if _, ok := f.store.User("sso:carol"); !ok {
		t.Fatalf("no account for a listed address: %v", f.store.Users())
	}
}

// Removing an entry revokes, rather than only refusing new accounts.
func TestOIDCAllowListRevokes(t *testing.T) {
	p := newFakeProvider(t)
	defer p.Close()
	f := oidcForge(t, p)
	h := f.routes()

	if w := login(t, f, h, p, nil); w.Code != http.StatusSeeOther {
		t.Fatalf("first login: %d", w.Code)
	}
	if _, ok := f.store.User("sso:carol"); !ok {
		t.Fatal("the first login provisioned nothing")
	}

	f.cfg.OIDC[0].Allow = []string{"@somewhere-else.example"}
	if w := login(t, f, h, p, nil); w.Code != http.StatusForbidden {
		t.Errorf("an account that exists still signed in after being removed: %d", w.Code)
	}
	if _, ok := f.store.User("sso:carol"); !ok {
		t.Error("the account was deleted rather than refused")
	}
}

// The other algorithm. A provider that signs with ES256 is verified by code
// that nothing else here runs, so it is exercised directly rather than through
// a second fake provider.
func TestES256IdTokens(t *testing.T) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	b64 := base64.RawURLEncoding.EncodeToString
	pub := key.Public().(*ecdsa.PublicKey)
	jwk := jwk{Kty: "EC", Kid: "e1", Use: "sig", Crv: "P-256",
		X: b64(pub.X.Bytes()), Y: b64(pub.Y.Bytes())}

	signed := []byte("the.token")
	sum := sha256.Sum256(signed)
	r, s, err := ecdsa.Sign(rand.Reader, key, sum[:])
	if err != nil {
		t.Fatal(err)
	}
	sig := make([]byte, 64)
	r.FillBytes(sig[:32])
	s.FillBytes(sig[32:])

	if err := verifySignature(jwk, "ES256", signed, sig); err != nil {
		t.Errorf("a good signature was refused: %v", err)
	}
	if err := verifySignature(jwk, "ES256", []byte("another.token"), sig); err == nil {
		t.Error("the signature verified against a different token")
	}
	// An algorithm the key cannot carry is refused rather than guessed at.
	if err := verifySignature(jwk, "RS256", signed, sig); err == nil {
		t.Error("an ec key verified an rs256 signature")
	}
	if err := verifySignature(jwk, "none", signed, nil); err == nil {
		t.Error("an unsigned token verified")
	}
	// A signature of the wrong size, and coordinates that are not on the
	// curve, both of which a hostile key set could hold.
	if err := verifySignature(jwk, "ES256", signed, sig[:63]); err == nil {
		t.Error("a short signature verified")
	}
	off := jwk
	off.X = b64([]byte{1, 2, 3})
	if err := verifySignature(off, "ES256", signed, sig); err == nil {
		t.Error("a point that is not on the curve verified")
	}
}