auth

func NewService(keys ...Key) (TokenService, error)

{
	if len(keys) == 0 {
		return nil, errors.New("auth: at least one key is required")
	}
	ids := make(map[string]struct{}, len(keys))
	for _, key := range keys {
		if key.ID == "" {
			return nil, errors.New("auth: key ID cannot be empty")
		}
		if _, exists := ids[key.ID]; exists {
			return nil, fmt.Errorf("auth: duplicate key ID %q", key.ID)
		}
		ids[key.ID] = struct{}{}
		if err := validateKey(key); err != nil {
			return nil, fmt.Errorf("auth: key %q: %w", key.ID, err)
		}
	}
	cloned := make([]Key, len(keys))
	for index, key := range keys {
		cloned[index] = cloneKey(key)
	}
	return &multiKeyService{keys: cloned}, nil
}

func cloneKey(key Key) Key

{
	cloned := key
	cloned.Secret = append([]byte(nil), key.Secret...)
	switch key.Algorithm {
	case AlgRS256:
		private := key.Private.(*rsa.PrivateKey)
		privateClone := &rsa.PrivateKey{
			PublicKey: rsa.PublicKey{N: new(big.Int).Set(private.N), E: private.E},
			D:         new(big.Int).Set(private.D),
			Primes:    make([]*big.Int, len(private.Primes)),
		}
		for index, prime := range private.Primes {
			privateClone.Primes[index] = new(big.Int).Set(prime)
		}
		privateClone.Precompute()
		public := key.Public.(*rsa.PublicKey)
		cloned.Private = privateClone
		cloned.Public = &rsa.PublicKey{N: new(big.Int).Set(public.N), E: public.E}
	case AlgES256:
		private := key.Private.(*ecdsa.PrivateKey)
		public := key.Public.(*ecdsa.PublicKey)
		cloned.Private = &ecdsa.PrivateKey{
			PublicKey: ecdsa.PublicKey{
				Curve: private.Curve,
				X:     new(big.Int).Set(private.X),
				Y:     new(big.Int).Set(private.Y),
			},
			D: new(big.Int).Set(private.D),
		}
		cloned.Public = &ecdsa.PublicKey{
			Curve: public.Curve,
			X:     new(big.Int).Set(public.X),
			Y:     new(big.Int).Set(public.Y),
		}
	case AlgEdDSA:
		cloned.Private = append(ed25519.PrivateKey(nil), key.Private.(ed25519.PrivateKey)...)
		cloned.Public = append(ed25519.PublicKey(nil), key.Public.(ed25519.PublicKey)...)
	}
	return cloned
}

func validateKey(key Key) error

{
	switch key.Algorithm {
	case AlgHMACSHA256:
		return ValidateHMACSecret(key.Secret)
	case AlgRS256:
		private, privateOK := key.Private.(*rsa.PrivateKey)
		public, publicOK := key.Public.(*rsa.PublicKey)
		if !privateOK || private == nil || !publicOK || public == nil {
			return errors.New("RS256 requires RSA private and public keys")
		}
		if err := private.Validate(); err != nil {
			return fmt.Errorf("invalid RSA private key: %w", err)
		}
		if private.N.BitLen() < 2048 {
			return errors.New("RS256 requires an RSA key of at least 2048 bits")
		}
		if private.PublicKey.E != public.E || private.PublicKey.N.Cmp(public.N) != 0 {
			return errors.New("RSA private and public keys do not match")
		}
	case AlgES256:
		private, privateOK := key.Private.(*ecdsa.PrivateKey)
		public, publicOK := key.Public.(*ecdsa.PublicKey)
		if !privateOK || private == nil || !publicOK || public == nil {
			return errors.New("ES256 requires ECDSA private and public keys")
		}
		if private.Curve != elliptic.P256() || public.Curve != elliptic.P256() {
			return errors.New("ES256 requires the P-256 curve")
		}
		if private.D == nil || private.D.Sign() <= 0 || private.D.Cmp(private.Curve.Params().N) >= 0 {
			return errors.New("invalid ECDSA private scalar")
		}
		derivedX, derivedY := private.Curve.ScalarBaseMult(private.D.Bytes())
		if private.X == nil || private.Y == nil || public.X == nil || public.Y == nil ||
			derivedX.Cmp(private.X) != 0 || derivedY.Cmp(private.Y) != 0 ||
			derivedX.Cmp(public.X) != 0 || derivedY.Cmp(public.Y) != 0 {
			return errors.New("ECDSA private and public keys do not match")
		}
	case AlgEdDSA:
		private, privateOK := key.Private.(ed25519.PrivateKey)
		public, publicOK := key.Public.(ed25519.PublicKey)
		if !privateOK || len(private) != ed25519.PrivateKeySize ||
			!publicOK || len(public) != ed25519.PublicKeySize {
			return errors.New("EdDSA requires Ed25519 private and public keys")
		}
		derived := private.Public().(ed25519.PublicKey)
		if !bytes.Equal(derived, public) {
			return errors.New("Ed25519 private and public keys do not match")
		}
	default:
		return fmt.Errorf("unsupported algorithm %q", key.Algorithm)
	}
	return nil
}

func signWithKey(claims Claims, key Key) (*SignedToken, error)

{
	data, err := json.Marshal(claims)
	if err != nil {
		return nil, err
	}

	var sig []byte
	switch key.Algorithm {
	case AlgHMACSHA256:
		sig, err = signHMAC(data, key.Secret)
	case AlgRS256:
		sig, err = signRSA(data, key.Private.(*rsa.PrivateKey))
	case AlgES256:
		sig, err = signECDSA(data, key.Private.(*ecdsa.PrivateKey))
	case AlgEdDSA:
		sig, err = signEd25519(data, key.Private.(ed25519.PrivateKey))
	default:
		return nil, fmt.Errorf("auth: unsupported algorithm: %s", key.Algorithm)
	}
	if err != nil {
		return nil, err
	}

	kid := base64.RawURLEncoding.EncodeToString([]byte(key.ID))
	payload := base64.RawURLEncoding.EncodeToString(data)
	signature := base64.RawURLEncoding.EncodeToString(sig)
	token := kid + "." + payload + "." + signature
	if len(token) > maxTokenSize {
		return nil, fmt.Errorf("auth: signed token exceeds %d bytes", maxTokenSize)
	}

	return &SignedToken{
		Token:   token,
		KeyID:   key.ID,
		Payload: data,
		Raw:     token,
	}, nil
}

func verifyWithKey[T Claims](token string, key Key) (T, error)

{
	var zero T

	parts, ok := splitToken(token)
	if !ok {
		return zero, ErrInvalidToken
	}

	kidBytes, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return zero, ErrInvalidToken
	}
	if string(kidBytes) != key.ID {
		return zero, ErrUnknownKeyID
	}

	payloadBytes, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return zero, ErrInvalidToken
	}

	sig, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return zero, ErrInvalidToken
	}

	switch key.Algorithm {
	case AlgRS256:
		if err := verifyRSA(payloadBytes, sig, key.Public.(*rsa.PublicKey)); err != nil {
			return zero, err
		}
	case AlgHMACSHA256:
		if !verifyHMAC(payloadBytes, sig, key.Secret) {
			return zero, ErrInvalidSignature
		}
	case AlgES256:
		if !verifyECDSA(payloadBytes, sig, key.Public.(*ecdsa.PublicKey)) {
			return zero, ErrInvalidSignature
		}
	case AlgEdDSA:
		if !verifyEd25519(payloadBytes, sig, key.Public.(ed25519.PublicKey)) {
			return zero, ErrInvalidSignature
		}
	default:
		return zero, fmt.Errorf("auth: unsupported algorithm: %s", key.Algorithm)
	}

	var claims T
	if err := json.Unmarshal(payloadBytes, &claims); err != nil {
		return zero, err
	}
	if c, ok := any(&claims).(Claims); ok {
		if err := c.Valid(); err != nil {
			return zero, err
		}
	}
	return claims, nil
}

func sign[T any](p T, secret []byte, alg Algorithm) (string, error)

{
	data, err := json.Marshal(p)
	if err != nil {
		return "", err
	}

	var sig []byte
	switch alg {
	case AlgHMACSHA256:
		sig = computeSig(data, secret)
	default:
		return "", fmt.Errorf("auth: unsupported algorithm: %s", alg)
	}

	token := base64.RawURLEncoding.EncodeToString(data) + "." + base64.RawURLEncoding.EncodeToString(sig)
	if len(token) > maxTokenSize {
		return "", fmt.Errorf("auth: signed token exceeds %d bytes", maxTokenSize)
	}
	return token, nil
}

func verify[T any](token string, secret []byte, alg Algorithm) (T, error)

{
	var zero T

	if len(token) == 0 || len(token) > maxTokenSize {
		return zero, ErrInvalidToken
	}
	payload, signature, ok := strings.Cut(token, ".")
	if !ok || payload == "" || signature == "" || strings.Contains(signature, ".") {
		return zero, ErrInvalidToken
	}

	payloadBytes, err := base64.RawURLEncoding.DecodeString(payload)
	if err != nil {
		return zero, ErrInvalidToken
	}

	sig, err := base64.RawURLEncoding.DecodeString(signature)
	if err != nil {
		return zero, ErrInvalidToken
	}

	var valid bool
	switch alg {
	case AlgHMACSHA256:
		expected := computeSig(payloadBytes, secret)
		valid = hmac.Equal(sig, expected)
	default:
		return zero, fmt.Errorf("auth: unsupported algorithm: %s", alg)
	}

	if !valid {
		return zero, ErrInvalidSignature
	}

	if err := json.Unmarshal(payloadBytes, &zero); err != nil {
		return zero, err
	}

	if p, ok := any(zero).(Claims); ok {
		if err := p.Valid(); err != nil {
			return zero, err
		}
	}

	return zero, nil
}

func splitToken(token string) ([3]string, bool)

{
	var parts [3]string
	if len(token) == 0 || len(token) > maxTokenSize {
		return parts, false
	}
	var ok bool
	parts[0], token, ok = strings.Cut(token, ".")
	if !ok {
		return parts, false
	}
	parts[1], parts[2], ok = strings.Cut(token, ".")
	if !ok || strings.Contains(parts[2], ".") ||
		parts[0] == "" || parts[1] == "" || parts[2] == "" {
		return parts, false
	}
	return parts, true
}

func TestSignTokenAndVerifyToken(t *testing.T)

{
	secret := testHMACSecret
	token, err := SignToken(Payload{Sub: "u1", Exp: time.Now().Add(time.Minute).Unix()}, secret)
	if err != nil {
		t.Fatalf("SignToken() error = %v", err)
	}
	payload, err := VerifyToken(token, secret)
	if err != nil {
		t.Fatalf("VerifyToken() error = %v", err)
	}
	if payload.Sub != "u1" {
		t.Fatalf("VerifyToken() subject = %q, want %q", payload.Sub, "u1")
	}
}

func TestVerifyTokenRejectsInvalidSignature(t *testing.T)

{
	secret := testHMACSecret
	token, err := SignToken(Payload{Sub: "u1", Exp: time.Now().Add(time.Minute).Unix()}, secret)
	if err != nil {
		t.Fatalf("SignToken() error = %v", err)
	}
	payload, _, ok := strings.Cut(token, ".")
	if !ok {
		t.Fatal("expected 2 parts")
	}
	tampered := payload + "." + base64.RawURLEncoding.EncodeToString([]byte("tampered"))
	_, err = VerifyToken(tampered, secret)
	if !errors.Is(err, ErrInvalidSignature) {
		t.Fatalf("VerifyToken() error = %v, want ErrInvalidSignature", err)
	}
}

func TestVerifyTokenRejectsExpiredToken(t *testing.T)

{
	secret := testHMACSecret
	token, err := SignToken(Payload{Sub: "u1", Exp: time.Now().Add(-time.Minute).Unix()}, secret)
	if err != nil {
		t.Fatalf("SignToken() error = %v", err)
	}
	_, err = VerifyToken(token, secret)
	if !errors.Is(err, ErrExpiredToken) {
		t.Fatalf("VerifyToken() error = %v, want ErrExpiredToken", err)
	}
}

func TestNewService_KeyRotation(t *testing.T)

{
	oldKey := Key{ID: "old", Secret: hmacSecret('o'), Algorithm: AlgHMACSHA256}
	newKey := Key{ID: "new", Secret: hmacSecret('n'), Algorithm: AlgHMACSHA256}
	svc := MustNewService(newKey, oldKey)
	token, err := svc.Sign(&StandardClaims{Sub: "user1", Exp: time.Now().Add(time.Hour).Unix()})
	if err != nil {
		t.Fatalf("Sign failed: %v", err)
	}
	if token.KeyID != "new" {
		t.Errorf("got key id %q, want %q", token.KeyID, "new")
	}
	claims, err := svc.Verify(token.Token)
	if err != nil {
		t.Fatalf("Verify failed: %v", err)
	}
	c := claims.(StandardClaims)
	if c.Sub != "user1" {
		t.Errorf("got sub %s, want user1", c.Sub)
	}
}

func TestNewServiceClonesHMACKeyMaterial(t *testing.T)

{
	secret := hmacSecret('k')
	service := MustNewService(Key{ID: "key", Secret: secret, Algorithm: AlgHMACSHA256})
	secret[0] ^= 0xff

	token, err := service.Sign(&StandardClaims{
		Sub: "user",
		Exp: time.Now().Add(time.Minute).Unix(),
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := service.Verify(token.Token); err != nil {
		t.Fatalf("Verify() after caller key mutation: %v", err)
	}
}

func TestNewService_KeyRotationOldKeyStillWorks(t *testing.T)

{
	oldKey := Key{ID: "old", Secret: hmacSecret('o'), Algorithm: AlgHMACSHA256}
	newKey := Key{ID: "new", Secret: hmacSecret('n'), Algorithm: AlgHMACSHA256}
	svcOld := MustNewService(oldKey)
	token, _ := svcOld.Sign(&StandardClaims{Sub: "user1", Exp: time.Now().Add(time.Hour).Unix()})
	svcNew := MustNewService(newKey, oldKey)
	claims, err := svcNew.Verify(token.Token)
	if err != nil {
		t.Fatalf("Verify with old key should still work: %v", err)
	}
	c := claims.(StandardClaims)
	if c.Sub != "user1" {
		t.Errorf("got sub %s, want user1", c.Sub)
	}
}

func TestNewService_WrongKey(t *testing.T)

{
	key1 := Key{ID: "k1", Secret: hmacSecret('1'), Algorithm: AlgHMACSHA256}
	key2 := Key{ID: "k2", Secret: hmacSecret('2'), Algorithm: AlgHMACSHA256}
	svc1 := MustNewService(key1)
	token, _ := svc1.Sign(&StandardClaims{Sub: "u", Exp: time.Now().Add(time.Hour).Unix()})
	svc2 := MustNewService(key2)
	_, err := svc2.Verify(token.Token)
	if err == nil {
		t.Error("expected error when verifying with wrong key")
	}
}

func TestNewService_RS256(t *testing.T)

{
	priv, _ := rsa.GenerateKey(rand.Reader, 2048)
	svc := MustNewService(Key{ID: "rsa1", Private: priv, Public: &priv.PublicKey, Algorithm: AlgRS256})
	token, err := svc.Sign(&StandardClaims{Sub: "user_rsa", Exp: time.Now().Add(time.Hour).Unix()})
	if err != nil {
		t.Fatalf("Sign with RSA failed: %v", err)
	}
	claims, err := svc.Verify(token.Token)
	if err != nil {
		t.Fatalf("Verify with RSA failed: %v", err)
	}
	c := claims.(StandardClaims)
	if c.Sub != "user_rsa" {
		t.Errorf("got sub %s, want user_rsa", c.Sub)
	}
}

func TestNewService_EdDSA(t *testing.T)

{
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKey failed: %v", err)
	}
	svc := MustNewService(Key{ID: "ed1", Private: priv, Public: pub, Algorithm: AlgEdDSA})
	token, err := svc.Sign(&StandardClaims{Sub: "user_ed", Exp: time.Now().Add(time.Hour).Unix()})
	if err != nil {
		t.Fatalf("Sign with EdDSA failed: %v", err)
	}
	claims, err := svc.Verify(token.Token)
	if err != nil {
		t.Fatalf("Verify with EdDSA failed: %v", err)
	}
	c := claims.(StandardClaims)
	if c.Sub != "user_ed" {
		t.Errorf("got sub %s, want user_ed", c.Sub)
	}
}

func TestNewService_ECDSA(t *testing.T)

{
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKey failed: %v", err)
	}
	svc := MustNewService(Key{ID: "ec1", Private: priv, Public: &priv.PublicKey, Algorithm: AlgES256})
	token, err := svc.Sign(&StandardClaims{Sub: "user_ec", Exp: time.Now().Add(time.Hour).Unix()})
	if err != nil {
		t.Fatalf("Sign with ECDSA failed: %v", err)
	}
	claims, err := svc.Verify(token.Token)
	if err != nil {
		t.Fatalf("Verify with ECDSA failed: %v", err)
	}
	c := claims.(StandardClaims)
	if c.Sub != "user_ec" {
		t.Errorf("got sub %s, want user_ec", c.Sub)
	}
}

func TestNewService_RejectsInvalidKeys(t *testing.T)

{
	tests := []Key{
		{ID: "short", Secret: []byte("short"), Algorithm: AlgHMACSHA256},
		{ID: "rsa", Algorithm: AlgRS256, Private: "wrong", Public: "wrong"},
		{ID: "unknown", Algorithm: Algorithm("unknown")},
	}
	for _, key := range tests {
		if _, err := NewService(key); err == nil {
			t.Fatalf("NewService() accepted key %#v", key)
		}
	}

	key := Key{ID: "same", Secret: testHMACSecret, Algorithm: AlgHMACSHA256}
	if _, err := NewService(key, key); err == nil {
		t.Fatal("NewService() accepted duplicate key IDs")
	}
}

func TestNewService_RejectsMismatchedAsymmetricKeys(t *testing.T)

{
	rsaPrivate, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	otherRSA, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	ecdsaPrivate, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	otherECDSA, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	_, edPrivate, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	otherEdPublic, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	p384Private, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	keys := []Key{
		{ID: "rsa", Private: rsaPrivate, Public: &otherRSA.PublicKey, Algorithm: AlgRS256},
		{ID: "ecdsa", Private: ecdsaPrivate, Public: &otherECDSA.PublicKey, Algorithm: AlgES256},
		{ID: "ed25519", Private: edPrivate, Public: otherEdPublic, Algorithm: AlgEdDSA},
		{ID: "p384", Private: p384Private, Public: &p384Private.PublicKey, Algorithm: AlgES256},
	}
	for _, key := range keys {
		if _, err := NewService(key); err == nil {
			t.Fatalf("NewService() accepted mismatched key %q", key.ID)
		}
	}
}

func TestVerifyRejectsOversizedTokens(t *testing.T)

{
	oversized := strings.Repeat(".", maxTokenSize+1)
	if _, err := VerifyToken(oversized, testHMACSecret); !errors.Is(err, ErrInvalidToken) {
		t.Fatalf("VerifyToken() error = %v, want ErrInvalidToken", err)
	}

	service, err := NewService(Key{
		ID:        "key",
		Secret:    testHMACSecret,
		Algorithm: AlgHMACSHA256,
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := service.Verify(oversized); err == nil {
		t.Fatal("TokenService.Verify() accepted an oversized token")
	}
}

func TestTokenServiceDoesNotCreateUnverifiableOversizedToken(t *testing.T)

{
	service, err := NewService(Key{
		ID:        "key",
		Secret:    testHMACSecret,
		Algorithm: AlgHMACSHA256,
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := service.Sign(largeClaims{Data: strings.Repeat("x", maxTokenSize)}); err == nil {
		t.Fatal("Sign() created a token larger than the verifier accepts")
	}
}