From dcaecc56f662109de334a9a3b790aec5bb772bf4 Mon Sep 17 00:00:00 2001
From: Nikos <nikos.sklikas@canonical.com>
Date: Wed, 16 Oct 2024 19:18:24 +0300
Subject: [PATCH] fix: make user code creation configurable

---
 config.go                           |  6 ++++++
 config_default.go                   | 23 +++++++++++++++++++++++
 fosite.go                           |  1 +
 handler/rfc8628/strategy_hmacsha.go |  3 ++-
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/config.go b/config.go
index 24c4151c..2d74753f 100644
--- a/config.go
+++ b/config.go
@@ -51,6 +51,12 @@ type DeviceAndUserCodeLifespanProvider interface {
 	GetDeviceAndUserCodeLifespan(ctx context.Context) time.Duration
 }
 
+// DeviceAndUserCodeLifespanProvider returns the provider for configuring the device and user code lifespan
+type UserCodeProvider interface {
+	GetUserCodeLength(ctx context.Context) int
+	GetUserCodeSymbols(ctx context.Context) []rune
+}
+
 // ScopeStrategyProvider returns the provider for configuring the scope strategy.
 type ScopeStrategyProvider interface {
 	// GetScopeStrategy returns the scope strategy.
diff --git a/config_default.go b/config_default.go
index a73154f0..eb4acce9 100644
--- a/config_default.go
+++ b/config_default.go
@@ -13,6 +13,7 @@ import (
 	"github.com/hashicorp/go-retryablehttp"
 
 	"github.com/ory/fosite/token/jwt"
+	"github.com/ory/x/randx"
 
 	"github.com/ory/fosite/i18n"
 )
@@ -229,6 +230,12 @@ type Config struct {
 
 	// IsPushedAuthorizeEnforced enforces pushed authorization request for /authorize
 	IsPushedAuthorizeEnforced bool
+
+	// UserCodeLength defines the length of the user_code
+	UserCodeLength int
+
+	// UserCodeSymbols defines the symbols that will be used to construct the user_code
+	UserCodeSymbols []rune
 }
 
 func (c *Config) GetGlobalSecret(ctx context.Context) ([]byte, error) {
@@ -540,3 +547,19 @@ func (c *Config) GetDeviceAuthTokenPollingInterval(ctx context.Context) time.Dur
 	}
 	return c.DeviceAuthTokenPollingInterval
 }
+
+// GetUserCodeLength returns configured user_code length
+func (c *Config) GetUserCodeLength(ctx context.Context) int {
+	if c.UserCodeLength == 0 {
+		return 8
+	}
+	return c.UserCodeLength
+}
+
+// GetDeviceAuthTokenPollingInterval returns configured user_code allowed symbols
+func (c *Config) GetUserCodeSymbols(ctx context.Context) []rune {
+	if c.UserCodeSymbols == nil {
+		return []rune(randx.AlphaUpper)
+	}
+	return c.UserCodeSymbols
+}
diff --git a/fosite.go b/fosite.go
index d5610129..90c30b4d 100644
--- a/fosite.go
+++ b/fosite.go
@@ -148,6 +148,7 @@ type Configurator interface {
 	RevocationHandlersProvider
 	UseLegacyErrorFormatProvider
 	DeviceEndpointHandlersProvider
+	UserCodeProvider
 	DeviceProvider
 }
 
diff --git a/handler/rfc8628/strategy_hmacsha.go b/handler/rfc8628/strategy_hmacsha.go
index 9f604096..6315bd92 100644
--- a/handler/rfc8628/strategy_hmacsha.go
+++ b/handler/rfc8628/strategy_hmacsha.go
@@ -103,6 +103,7 @@ type DefaultDeviceStrategy struct {
 	Config interface {
 		fosite.DeviceProvider
 		fosite.DeviceAndUserCodeLifespanProvider
+		fosite.UserCodeProvider
 	}
 }
 
@@ -110,7 +111,7 @@ var _ RFC8628CodeStrategy = (*DefaultDeviceStrategy)(nil)
 
 // GenerateUserCode generates a user_code
 func (h *DefaultDeviceStrategy) GenerateUserCode(ctx context.Context) (string, string, error) {
-	seq, err := randx.RuneSequence(8, []rune(randx.AlphaUpper))
+	seq, err := randx.RuneSequence(h.Config.GetUserCodeLength(ctx), h.Config.GetUserCodeSymbols(ctx))
 	if err != nil {
 		return "", "", err
 	}