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 ?