-

2025-08-23 16:22:46 -04:00
parent 0c50b05324
commit 3e02ef57b9
9 changed files with 196 additions and 18 deletions
--- a/internal/crypto/token.go
+++ b/internal/crypto/token.go
@ -34,6 +34,11 @@ func NewTokenGenerator(hmacKey string) *TokenGenerator {

 // GenerateSecureToken generates a cryptographically secure random token
 func (tg *TokenGenerator) GenerateSecureToken() (string, error) {
+	return tg.GenerateSecureTokenWithPrefix("", "")
+}
+
+// GenerateSecureTokenWithPrefix generates a cryptographically secure random token with custom prefix
+func (tg *TokenGenerator) GenerateSecureTokenWithPrefix(appPrefix string, tokenType string) (string, error) {
 	// Generate random bytes
 	tokenBytes := make([]byte, TokenLength)
 	if _, err := rand.Read(tokenBytes); err != nil {
@ -43,8 +48,21 @@ func (tg *TokenGenerator) GenerateSecureToken() (string, error) {
 	// Encode to base64 for safe transmission
 	tokenData := base64.URLEncoding.EncodeToString(tokenBytes)
 	
-	// Add prefix for identification
-	token := TokenPrefix + tokenData
+	// Build prefix based on application and token type
+	var prefix string
+	if appPrefix != "" {
+		// Use custom application prefix
+		if tokenType == "user" {
+			prefix = appPrefix + "UT-" // User Token
+		} else {
+			prefix = appPrefix + "T-"  // Static Token
+		}
+	} else {
+		// Use default prefix
+		prefix = TokenPrefix
+	}
+	
+	token := prefix + tokenData

 	return token, nil
 }
@ -103,12 +121,52 @@ func ExtractTokenFromHeader(authHeader string) string {

 // IsValidTokenFormat checks if a token has the expected format
 func IsValidTokenFormat(token string) bool {
-	if !strings.HasPrefix(token, TokenPrefix) {
-		return false
+	return IsValidTokenFormatWithPrefix(token, "")
+}
+
+// IsValidTokenFormatWithPrefix checks if a token has the expected format with custom prefix
+func IsValidTokenFormatWithPrefix(token string, expectedPrefix string) bool {
+	var prefix string
+	if expectedPrefix != "" {
+		prefix = expectedPrefix
+	} else {
+		prefix = TokenPrefix
+	}
+	
+	if !strings.HasPrefix(token, prefix) {
+		// If expected prefix doesn't match, check if it's a valid token with any custom prefix
+		if expectedPrefix == "" {
+			// Check for custom prefix pattern: 2-4 uppercase letters + "T-" or "UT-"
+			if len(token) < 6 { // minimum: "ABT-" + some data
+				return false
+			}
+			
+			// Look for T- or UT- suffix in the first part
+			dashIndex := strings.Index(token, "-")
+			if dashIndex < 2 || dashIndex > 6 { // 2-4 chars + "T" or "UT"
+				// Not a custom prefix, check default
+				if !strings.HasPrefix(token, TokenPrefix) {
+					return false
+				}
+				prefix = TokenPrefix
+			} else {
+				prefixPart := token[:dashIndex+1]
+				if !strings.HasSuffix(prefixPart, "T-") && !strings.HasSuffix(prefixPart, "UT-") {
+					if !strings.HasPrefix(token, TokenPrefix) {
+						return false
+					}
+					prefix = TokenPrefix
+				} else {
+					prefix = prefixPart
+				}
+			}
+		} else {
+			return false
+		}
 	}
 	
 	// Remove prefix and check if remaining part is valid base64
-	tokenData := strings.TrimPrefix(token, TokenPrefix)
+	tokenData := strings.TrimPrefix(token, prefix)
 	if len(tokenData) == 0 {
 		return false
 	}
@ -128,8 +186,13 @@ type TokenInfo struct {

 // GenerateTokenWithInfo generates a complete token with hash and signature
 func (tg *TokenGenerator) GenerateTokenWithInfo() (*TokenInfo, error) {
+	return tg.GenerateTokenWithInfoAndPrefix("", "")
+}
+
+// GenerateTokenWithInfoAndPrefix generates a complete token with hash, signature, and custom prefix
+func (tg *TokenGenerator) GenerateTokenWithInfoAndPrefix(appPrefix string, tokenType string) (*TokenInfo, error) {
 	// Generate the token
-	token, err := tg.GenerateSecureToken()
+	token, err := tg.GenerateSecureTokenWithPrefix(appPrefix, tokenType)
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate token: %w", err)
 	}
--- a/internal/domain/models.go
+++ b/internal/domain/models.go
@ -44,6 +44,7 @@ type Application struct {
 	Type                 []ApplicationType `json:"type" validate:"required,min=1,dive,oneof=static user" db:"type"`
 	CallbackURL          string            `json:"callback_url" validate:"required,url,max=500" db:"callback_url"`
 	HMACKey              string            `json:"hmac_key" validate:"required,min=1,max=255" db:"hmac_key"`
+	TokenPrefix          string            `json:"token_prefix" validate:"omitempty,min=2,max=4,uppercase" db:"token_prefix"`
 	TokenRenewalDuration Duration          `json:"token_renewal_duration" validate:"required,min=1" db:"token_renewal_duration"`
 	MaxTokenDuration     Duration          `json:"max_token_duration" validate:"required,min=1" db:"max_token_duration"`
 	Owner                Owner             `json:"owner" validate:"required"`
@ -158,6 +159,7 @@ type CreateApplicationRequest struct {
 	AppLink              string            `json:"app_link" validate:"required,url,max=500"`
 	Type                 []ApplicationType `json:"type" validate:"required,min=1,dive,oneof=static user"`
 	CallbackURL          string            `json:"callback_url" validate:"required,url,max=500"`
+	TokenPrefix          string            `json:"token_prefix" validate:"omitempty,min=2,max=4,uppercase"`
 	TokenRenewalDuration Duration          `json:"token_renewal_duration" validate:"required,min=1"`
 	MaxTokenDuration     Duration          `json:"max_token_duration" validate:"required,min=1"`
 	Owner                Owner             `json:"owner" validate:"required"`
@ -169,6 +171,7 @@ type UpdateApplicationRequest struct {
 	Type                 *[]ApplicationType `json:"type,omitempty" validate:"omitempty,min=1,dive,oneof=static user"`
 	CallbackURL          *string            `json:"callback_url,omitempty" validate:"omitempty,url,max=500"`
 	HMACKey              *string            `json:"hmac_key,omitempty" validate:"omitempty,min=1,max=255"`
+	TokenPrefix          *string            `json:"token_prefix,omitempty" validate:"omitempty,min=2,max=4,uppercase"`
 	TokenRenewalDuration *Duration          `json:"token_renewal_duration,omitempty" validate:"omitempty,min=1"`
 	MaxTokenDuration     *Duration          `json:"max_token_duration,omitempty" validate:"omitempty,min=1"`
 	Owner                *Owner             `json:"owner,omitempty" validate:"omitempty"`
--- a/internal/repository/postgres/application_repository.go
+++ b/internal/repository/postgres/application_repository.go
@ -26,11 +26,11 @@ func NewApplicationRepository(db repository.DatabaseProvider) repository.Applica
 func (r *ApplicationRepository) Create(ctx context.Context, app *domain.Application) error {
 	query := `
 		INSERT INTO applications (
-			app_id, app_link, type, callback_url, hmac_key, 
+			app_id, app_link, type, callback_url, hmac_key, token_prefix,
 			token_renewal_duration, max_token_duration, 
 			owner_type, owner_name, owner_owner,
 			created_at, updated_at
-		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
+		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
 	`

 	db := r.db.GetDB().(*sql.DB)
@ -48,6 +48,7 @@ func (r *ApplicationRepository) Create(ctx context.Context, app *domain.Applicat
 		pq.Array(typeStrings),
 		app.CallbackURL,
 		app.HMACKey,
+		app.TokenPrefix,
 		app.TokenRenewalDuration.Duration.Nanoseconds(),
 		app.MaxTokenDuration.Duration.Nanoseconds(),
 		string(app.Owner.Type),
@ -73,7 +74,7 @@ func (r *ApplicationRepository) Create(ctx context.Context, app *domain.Applicat
 // GetByID retrieves an application by its ID
 func (r *ApplicationRepository) GetByID(ctx context.Context, appID string) (*domain.Application, error) {
 	query := `
-		SELECT app_id, app_link, type, callback_url, hmac_key,
+		SELECT app_id, app_link, type, callback_url, hmac_key, token_prefix,
 		       token_renewal_duration, max_token_duration,
 		       owner_type, owner_name, owner_owner,
 		       created_at, updated_at
@ -95,6 +96,7 @@ func (r *ApplicationRepository) GetByID(ctx context.Context, appID string) (*dom
 		&typeStrings,
 		&app.CallbackURL,
 		&app.HMACKey,
+		&app.TokenPrefix,
 		&tokenRenewalNanos,
 		&maxTokenNanos,
 		&ownerType,
@ -130,7 +132,7 @@ func (r *ApplicationRepository) GetByID(ctx context.Context, appID string) (*dom
 // List retrieves applications with pagination
 func (r *ApplicationRepository) List(ctx context.Context, limit, offset int) ([]*domain.Application, error) {
 	query := `
-		SELECT app_id, app_link, type, callback_url, hmac_key,
+		SELECT app_id, app_link, type, callback_url, hmac_key, token_prefix,
 		       token_renewal_duration, max_token_duration,
 		       owner_type, owner_name, owner_owner,
 		       created_at, updated_at
@ -160,6 +162,7 @@ func (r *ApplicationRepository) List(ctx context.Context, limit, offset int) ([]
 			&typeStrings,
 			&app.CallbackURL,
 			&app.HMACKey,
+			&app.TokenPrefix,
 			&tokenRenewalNanos,
 			&maxTokenNanos,
 			&ownerType,
@ -231,6 +234,12 @@ func (r *ApplicationRepository) Update(ctx context.Context, appID string, update
 		argIndex++
 	}

+	if updates.TokenPrefix != nil {
+		setParts = append(setParts, fmt.Sprintf("token_prefix = $%d", argIndex))
+		args = append(args, *updates.TokenPrefix)
+		argIndex++
+	}
+
 	if updates.TokenRenewalDuration != nil {
 		setParts = append(setParts, fmt.Sprintf("token_renewal_duration = $%d", argIndex))
 		args = append(args, updates.TokenRenewalDuration.Duration.Nanoseconds())
--- a/internal/services/application_service.go
+++ b/internal/services/application_service.go
@ -37,6 +37,7 @@ func (s *applicationService) Create(ctx context.Context, req *domain.CreateAppli
 		Type:                 req.Type,
 		CallbackURL:          req.CallbackURL,
 		HMACKey:              generateHMACKey(), // TODO: Use proper key generation
+		TokenPrefix:          req.TokenPrefix,
 		TokenRenewalDuration: req.TokenRenewalDuration,
 		MaxTokenDuration:     req.MaxTokenDuration,
 		Owner:                req.Owner,
--- a/internal/services/token_service.go
+++ b/internal/services/token_service.go
@ -3,6 +3,7 @@ package services
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"

 	"github.com/google/uuid"
@ -72,8 +73,8 @@ func (s *tokenService) CreateStaticToken(ctx context.Context, req *domain.Create
 		return nil, fmt.Errorf("some requested permissions are invalid")
 	}

-	// Generate secure token
-	tokenInfo, err := s.tokenGen.GenerateTokenWithInfo()
+	// Generate secure token with custom prefix
+	tokenInfo, err := s.tokenGen.GenerateTokenWithInfoAndPrefix(app.TokenPrefix, "static")
 	if err != nil {
 		s.logger.Error("Failed to generate secure token", zap.Error(err))
 		return nil, fmt.Errorf("failed to generate token: %w", err)
@ -239,12 +240,21 @@ func (s *tokenService) GenerateUserToken(ctx context.Context, appID, userID stri
 	}

 	// Generate JWT token using JWT manager
-	tokenString, err := s.jwtManager.GenerateToken(userToken)
+	jwtTokenString, err := s.jwtManager.GenerateToken(userToken)
 	if err != nil {
 		s.logger.Error("Failed to generate JWT token", zap.Error(err))
 		return "", fmt.Errorf("failed to generate token: %w", err)
 	}

+	// Add custom prefix wrapper for user tokens if application has one
+	var finalToken string
+	if app.TokenPrefix != "" {
+		// For user JWT tokens, we wrap the JWT with custom prefix
+		finalToken = app.TokenPrefix + "UT-" + jwtTokenString
+	} else {
+		finalToken = jwtTokenString
+	}
+
 	s.logger.Info("User token generated successfully", 
 		zap.String("app_id", appID),
 		zap.String("user_id", userID),
@ -252,7 +262,7 @@ func (s *tokenService) GenerateUserToken(ctx context.Context, appID, userID stri
 		zap.Time("expires_at", userToken.ExpiresAt),
 		zap.Time("max_valid_at", userToken.MaxValidAt))

-	return tokenString, nil
+	return finalToken, nil
 }

 // VerifyToken verifies a token and returns verification response
@ -389,8 +399,27 @@ func (s *tokenService) verifyStaticToken(ctx context.Context, req *domain.Verify
 func (s *tokenService) verifyUserToken(ctx context.Context, req *domain.VerifyRequest, app *domain.Application) (*domain.VerifyResponse, error) {
 	s.logger.Debug("Verifying user token", zap.String("app_id", req.AppID))

+	// Extract JWT token from potentially prefixed format
+	jwtToken := req.Token
+	if app.TokenPrefix != "" {
+		expectedPrefix := app.TokenPrefix + "UT-"
+		if strings.HasPrefix(req.Token, expectedPrefix) {
+			jwtToken = strings.TrimPrefix(req.Token, expectedPrefix)
+		} else {
+			// Token doesn't have expected prefix
+			s.logger.Warn("User token missing expected prefix", 
+				zap.String("app_id", req.AppID),
+				zap.String("expected_prefix", expectedPrefix))
+			return &domain.VerifyResponse{
+				Valid:     false,
+				Permitted: false,
+				Error:     "Invalid token format",
+			}, nil
+		}
+	}
+
 	// Check if token is revoked first
-	isRevoked, err := s.jwtManager.IsTokenRevoked(req.Token)
+	isRevoked, err := s.jwtManager.IsTokenRevoked(jwtToken)
 	if err != nil {
 		s.logger.Error("Failed to check token revocation status", zap.Error(err))
 		return &domain.VerifyResponse{
@ -410,7 +439,7 @@ func (s *tokenService) verifyUserToken(ctx context.Context, req *domain.VerifyRe
 	}

 	// Validate JWT token
-	claims, err := s.jwtManager.ValidateToken(req.Token)
+	claims, err := s.jwtManager.ValidateToken(jwtToken)
 	if err != nil {
 		s.logger.Warn("JWT token validation failed", zap.Error(err), zap.String("app_id", req.AppID))
 		return &domain.VerifyResponse{