auth

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

	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 := splitToken(token)
	if len(parts) != 3 {
		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)
	}

	return base64.RawURLEncoding.EncodeToString(data) + "." + base64.RawURLEncoding.EncodeToString(sig), nil
}

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

{
	var zero T

	idx := -1
	for i := 0; i < len(token); i++ {
		if token[i] == '.' {
			idx = i
			break
		}
	}
	if idx < 0 {
		return zero, ErrInvalidToken
	}

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

	sig, err := base64.RawURLEncoding.DecodeString(token[idx+1:])
	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) []string

{
	var parts []string
	start := 0
	for i := 0; i < len(token); i++ {
		if token[i] == '.' {
			parts = append(parts, token[start:i])
			start = i + 1
		}
	}
	if start < len(token) {
		parts = append(parts, token[start:])
	}
	return parts
}

func TestSignTokenAndVerifyToken(t *testing.T)

{
	secret := []byte("s3cr3t")
	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 := []byte("s3cr3t")
	token, err := SignToken(Payload{Sub: "u1", Exp: time.Now().Add(time.Minute).Unix()}, secret)
	if err != nil {
		t.Fatalf("SignToken() error = %v", err)
	}
	parts := splitToken(token)
	if len(parts) != 2 {
		t.Fatal("expected 2 parts")
	}
	tampered := parts[0] + "." + 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 := []byte("s3cr3t")
	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: []byte("old-secret"), Algorithm: AlgHMACSHA256}
	newKey := Key{ID: "new", Secret: []byte("new-secret"), Algorithm: AlgHMACSHA256}
	svc := NewService(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 TestNewService_KeyRotationOldKeyStillWorks(t *testing.T)

{
	oldKey := Key{ID: "old", Secret: []byte("old-secret"), Algorithm: AlgHMACSHA256}
	newKey := Key{ID: "new", Secret: []byte("new-secret"), Algorithm: AlgHMACSHA256}
	svcOld := NewService(oldKey)
	token, _ := svcOld.Sign(&StandardClaims{Sub: "user1", Exp: time.Now().Add(time.Hour).Unix()})
	svcNew := NewService(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: []byte("secret1"), Algorithm: AlgHMACSHA256}
	key2 := Key{ID: "k2", Secret: []byte("secret2"), Algorithm: AlgHMACSHA256}
	svc1 := NewService(key1)
	token, _ := svc1.Sign(&StandardClaims{Sub: "u", Exp: time.Now().Add(time.Hour).Unix()})
	svc2 := NewService(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 := NewService(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 := NewService(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 := NewService(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)
	}
}