diff --git a/internal/crypto/token.go b/internal/crypto/token.go index a631525..9fe3d59 100644 --- 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) } diff --git a/internal/domain/models.go b/internal/domain/models.go index 9141f4b..59a458f 100644 --- 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"` diff --git a/internal/repository/postgres/application_repository.go b/internal/repository/postgres/application_repository.go index eaddf61..b072aa3 100644 --- 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()) diff --git a/internal/services/application_service.go b/internal/services/application_service.go index 0afe84a..ad9da06 100644 --- 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, diff --git a/internal/services/token_service.go b/internal/services/token_service.go index 2378431..c26b3b2 100644 --- 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{ diff --git a/kms-frontend/src/components/Applications.tsx b/kms-frontend/src/components/Applications.tsx index e4cd35a..2b87fb3 100644 --- a/kms-frontend/src/components/Applications.tsx +++ b/kms-frontend/src/components/Applications.tsx @@ -67,6 +67,7 @@ const Applications: React.FC = () => { app_link: app.app_link, type: app.type, callback_url: app.callback_url, + token_prefix: app.token_prefix, token_renewal_duration: formatDuration(app.token_renewal_duration), max_token_duration: formatDuration(app.max_token_duration), owner_type: app.owner.type, @@ -94,6 +95,7 @@ const Applications: React.FC = () => { app_link: values.app_link, type: values.type, callback_url: values.callback_url, + token_prefix: values.token_prefix, token_renewal_duration: values.token_renewal_duration, max_token_duration: values.max_token_duration, owner: { @@ -165,6 +167,12 @@ const Applications: React.FC = () => { ), }, + { + title: 'Token Prefix', + dataIndex: 'token_prefix', + key: 'token_prefix', + render: (prefix: string) => prefix ? {prefix} : Default, + }, { title: 'Owner', dataIndex: 'owner', @@ -307,6 +315,28 @@ const Applications: React.FC = () => { + + { + const value = e.target.value.toUpperCase(); + form.setFieldValue('token_prefix', value); + }} + /> + + { - Token Renewal Duration: -
{formatDuration(selectedApp.token_renewal_duration)}
+ Token Prefix: +
+ {selectedApp.token_prefix ? ( + <> + {selectedApp.token_prefix} + + (Static: {selectedApp.token_prefix}T-, User: {selectedApp.token_prefix}UT-) + + + ) : ( + Default (kms_) + )} +
+ + Token Renewal Duration: +
{formatDuration(selectedApp.token_renewal_duration)}
+ Max Token Duration:
{formatDuration(selectedApp.max_token_duration)}
diff --git a/kms-frontend/src/services/apiService.ts b/kms-frontend/src/services/apiService.ts index 1b7f2e5..785acf6 100644 --- a/kms-frontend/src/services/apiService.ts +++ b/kms-frontend/src/services/apiService.ts @@ -7,6 +7,7 @@ export interface Application { type: string[]; callback_url: string; hmac_key: string; + token_prefix?: string; token_renewal_duration: number; max_token_duration: number; owner: { @@ -36,6 +37,7 @@ export interface CreateApplicationRequest { app_link: string; type: string[]; callback_url: string; + token_prefix?: string; token_renewal_duration: string; max_token_duration: string; owner: { diff --git a/migrations/003_add_token_prefix.down.sql b/migrations/003_add_token_prefix.down.sql new file mode 100644 index 0000000..0d7f372 --- /dev/null +++ b/migrations/003_add_token_prefix.down.sql @@ -0,0 +1,11 @@ +-- Migration: 003_add_token_prefix (down) +-- Remove token prefix field from applications table + +-- Drop constraint first +ALTER TABLE applications DROP CONSTRAINT IF EXISTS chk_token_prefix_format; + +-- Drop index +DROP INDEX IF EXISTS idx_applications_token_prefix; + +-- Drop the prefix column +ALTER TABLE applications DROP COLUMN IF EXISTS token_prefix; \ No newline at end of file diff --git a/migrations/003_add_token_prefix.up.sql b/migrations/003_add_token_prefix.up.sql new file mode 100644 index 0000000..ca0a10c --- /dev/null +++ b/migrations/003_add_token_prefix.up.sql @@ -0,0 +1,15 @@ +-- Migration: 003_add_token_prefix +-- Add token prefix field to applications table + +-- Add prefix field to applications table +ALTER TABLE applications ADD COLUMN token_prefix VARCHAR(10) DEFAULT '' NOT NULL; + +-- Add check constraint to ensure prefix is 2-4 uppercase letters +ALTER TABLE applications ADD CONSTRAINT chk_token_prefix_format + CHECK (token_prefix ~ '^[A-Z]{2,4}$' OR token_prefix = ''); + +-- Create index for prefix field +CREATE INDEX idx_applications_token_prefix ON applications(token_prefix); + +-- Update existing applications with empty prefix (they will use the default "kms_" prefix) +-- Applications can later be updated to have custom prefixes \ No newline at end of file