fix: handle the user code generation duplication

generate-mocks.sh

@@ -7,6 +7,7 @@ mockgen -package internal -destination internal/oauth2_storage.go github.com/ory
 mockgen -package internal -destination internal/oauth2_strategy.go github.com/ory/fosite/handler/oauth2 CoreStrategy
 mockgen -package internal -destination internal/authorize_code_storage.go github.com/ory/fosite/handler/oauth2 AuthorizeCodeStorage
 mockgen -package internal -destination internal/device_code_storage.go github.com/ory/fosite/handler/rfc8628 DeviceCodeStorage
+mockgen -package internal -destination internal/rfc8628_core_storage.go github.com/ory/fosite/handler/rfc8628 RFC8628CoreStorage
 mockgen -package internal -destination internal/oauth2_auth_jwt_storage.go github.com/ory/fosite/handler/rfc7523 RFC7523KeyStorage
 mockgen -package internal -destination internal/access_token_storage.go github.com/ory/fosite/handler/oauth2 AccessTokenStorage
 mockgen -package internal -destination internal/refresh_token_strategy.go github.com/ory/fosite/handler/oauth2 RefreshTokenStorage
@@ -17,6 +18,7 @@ mockgen -package internal -destination internal/openid_id_token_storage.go githu
 mockgen -package internal -destination internal/access_token_strategy.go github.com/ory/fosite/handler/oauth2 AccessTokenStrategy
 mockgen -package internal -destination internal/refresh_token_strategy.go github.com/ory/fosite/handler/oauth2 RefreshTokenStrategy
 mockgen -package internal -destination internal/authorize_code_strategy.go github.com/ory/fosite/handler/oauth2 AuthorizeCodeStrategy
+mockgen -package internal -destination internal/rfc8628_code_strategy.go github.com/ory/fosite/handler/rfc8628 RFC8628CodeStrategy
 mockgen -package internal -destination internal/device_code_rate_limit_strategy.go github.com/ory/fosite/handler/rfc8628 DeviceRateLimitStrategy
 mockgen -package internal -destination internal/id_token_strategy.go github.com/ory/fosite/handler/openid OpenIDConnectTokenStrategy
 mockgen -package internal -destination internal/pkce_storage_strategy.go github.com/ory/fosite/handler/pkce PKCERequestStorage

generate.go

@@ -10,6 +10,7 @@ package fosite
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/oauth2_strategy.go github.com/ory/fosite/handler/oauth2 CoreStrategy
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/authorize_code_storage.go github.com/ory/fosite/handler/oauth2 AuthorizeCodeStorage
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/device_code_storage.go github.com/ory/fosite/handler/rfc8628 DeviceCodeStorage
+//go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/rfc8628_core_storage.go github.com/ory/fosite/handler/rfc8628 RFC8628CoreStorage
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/oauth2_auth_jwt_storage.go github.com/ory/fosite/handler/rfc7523 RFC7523KeyStorage
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/access_token_storage.go github.com/ory/fosite/handler/oauth2 AccessTokenStorage
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/refresh_token_strategy.go github.com/ory/fosite/handler/oauth2 RefreshTokenStorage
@@ -21,6 +22,7 @@ package fosite
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/refresh_token_strategy.go github.com/ory/fosite/handler/oauth2 RefreshTokenStrategy
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/authorize_code_strategy.go github.com/ory/fosite/handler/oauth2 AuthorizeCodeStrategy
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/device_code_rate_limit_strategy.go github.com/ory/fosite/handler/rfc8628 DeviceRateLimitStrategy
+//go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/rfc8628_code_strategy.go github.com/ory/fosite/handler/rfc8628 RFC8628CodeStrategy
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/id_token_strategy.go github.com/ory/fosite/handler/openid OpenIDConnectTokenStrategy
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/pkce_storage_strategy.go github.com/ory/fosite/handler/pkce PKCERequestStorage
 //go:generate go run github.com/golang/mock/mockgen -package internal -destination internal/authorize_handler.go github.com/ory/fosite AuthorizeEndpointHandler

handler/rfc8628/auth_handler.go

@@ -5,13 +5,16 @@ package rfc8628
 
 import (
 	"context"
+	"fmt"
 	"time"
 
-	"github.com/ory/x/errorsx"
-
 	"github.com/ory/fosite"
+	"github.com/ory/x/errorsx"
 )
 
+// MaxAttempts for retrying the generation of user codes.
+const MaxAttempts = 3
+
 // DeviceAuthHandler is a response handler for the Device Authorisation Grant as
 // defined in https://tools.ietf.org/html/rfc8628#section-3.1
 type DeviceAuthHandler struct {
@@ -25,25 +28,18 @@ type DeviceAuthHandler struct {
 
 // HandleDeviceEndpointRequest implements https://tools.ietf.org/html/rfc8628#section-3.1
 func (d *DeviceAuthHandler) HandleDeviceEndpointRequest(ctx context.Context, dar fosite.DeviceRequester, resp fosite.DeviceResponder) error {
-	deviceCode, deviceCodeSignature, err := d.Strategy.GenerateDeviceCode(ctx)
-	if err != nil {
-		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
-	}
+	var err error
 
-	userCode, userCodeSignature, err := d.Strategy.GenerateUserCode(ctx)
+	var deviceCode string
+	deviceCode, err = d.handleDeviceCode(ctx, dar)
 	if err != nil {
-		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
+		return err
 	}
 
-	// Store the User Code session (this has no real data other that the user and device code), can be converted into a 'full' session after user auth
-	dar.GetSession().SetExpiresAt(fosite.DeviceCode, time.Now().UTC().Add(d.Config.GetDeviceAndUserCodeLifespan(ctx)))
-	if err := d.Storage.CreateDeviceCodeSession(ctx, deviceCodeSignature, dar.Sanitize(nil)); err != nil {
-		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
-	}
-
-	dar.GetSession().SetExpiresAt(fosite.UserCode, time.Now().UTC().Add(d.Config.GetDeviceAndUserCodeLifespan(ctx)).Round(time.Second))
-	if err := d.Storage.CreateUserCodeSession(ctx, userCodeSignature, dar.Sanitize(nil)); err != nil {
-		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
+	var userCode string
+	userCode, err = d.handleUserCode(ctx, dar)
+	if err != nil {
+		return err
 	}
 
 	// Populate the response fields
@@ -55,3 +51,42 @@ func (d *DeviceAuthHandler) HandleDeviceEndpointRequest(ctx context.Context, dar
 	resp.SetInterval(int(d.Config.GetDeviceAuthTokenPollingInterval(ctx).Seconds()))
 	return nil
 }
+
+func (d *DeviceAuthHandler) handleDeviceCode(ctx context.Context, dar fosite.DeviceRequester) (string, error) {
+	code, signature, err := d.Strategy.GenerateDeviceCode(ctx)
+	if err != nil {
+		return "", errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
+	}
+
+	dar.GetSession().SetExpiresAt(fosite.DeviceCode, time.Now().UTC().Add(d.Config.GetDeviceAndUserCodeLifespan(ctx)))
+	if err = d.Storage.CreateDeviceCodeSession(ctx, signature, dar.Sanitize(nil)); err != nil {
+		return "", errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
+	}
+
+	return code, nil
+}
+
+func (d *DeviceAuthHandler) handleUserCode(ctx context.Context, dar fosite.DeviceRequester) (string, error) {
+	var err error
+	var userCode, signature string
+	// Note: the retries are added here because we need to ensure uniqueness of user codes.
+	// The chances of duplicates should however be diminishing, because they are the same
+	// chance an attacker will be able to hit a valid code with few guesses. However, as
+	// used codes will probably still be around for some time before they get cleaned,
+	// the chances of hitting a duplicate here can be higher.
+	// Three retries should be plenty, as otherwise the entropy is definitely off.
+	for i := 0; i < MaxAttempts; i++ {
+		userCode, signature, err = d.Strategy.GenerateUserCode(ctx)
+		if err != nil {
+			return "", err
+		}
+
+		dar.GetSession().SetExpiresAt(fosite.UserCode, time.Now().UTC().Add(d.Config.GetDeviceAndUserCodeLifespan(ctx)).Round(time.Second))
+		if err = d.Storage.CreateUserCodeSession(ctx, signature, dar.Sanitize(nil)); err == nil {
+			return userCode, nil
+		}
+	}
+
+	errMsg := fmt.Sprintf("Exceeded user-code generation max attempts %v: %s", MaxAttempts, err.Error())
+	return "", errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(errMsg))
+}

handler/rfc8628/auth_handler_test.go

@@ -5,15 +5,20 @@ package rfc8628_test
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"testing"
 	"time"
 
-	"github.com/ory/fosite/storage"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
+	"github.com/golang/mock/gomock"
+
+	"github.com/ory/fosite/internal"
 
 	"github.com/ory/fosite"
 	"github.com/ory/fosite/handler/rfc8628"
+	"github.com/ory/fosite/storage"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func Test_HandleDeviceEndpointRequest(t *testing.T) {
@@ -52,3 +57,155 @@ func Test_HandleDeviceEndpointRequest(t *testing.T) {
 	assert.Contains(t, resp.GetDeviceCode(), ".")
 	assert.Equal(t, resp.GetVerificationURI(), "www.test.com")
 }
+
+func Test_HandleDeviceEndpointRequestWithRetry(t *testing.T) {
+	var mockRFC8628CoreStorage *internal.MockRFC8628CoreStorage
+	var mockRFC8628CodeStrategy *internal.MockRFC8628CodeStrategy
+
+	ctx := context.Background()
+	req := &fosite.DeviceRequest{
+		Request: fosite.Request{
+			Client: &fosite.DefaultClient{
+				Audience: []string{"https://www.ory.sh/api"},
+			},
+			Session: &fosite.DefaultSession{},
+		},
+	}
+
+	testCases := []struct {
+		description string
+		setup       func()
+		check       func(t *testing.T, resp *fosite.DeviceResponse)
+		expectError error
+	}{
+		{
+			description: "should pass when generating a unique user code at the first attempt",
+			setup: func() {
+				mockRFC8628CodeStrategy.
+					EXPECT().
+					GenerateDeviceCode(ctx).
+					Return("deviceCode", "signature", nil)
+				mockRFC8628CoreStorage.
+					EXPECT().
+					CreateDeviceCodeSession(ctx, "signature", gomock.Any()).
+					Return(nil)
+				mockRFC8628CodeStrategy.
+					EXPECT().
+					GenerateUserCode(ctx).
+					Return("userCode", "signature", nil).
+					Times(1)
+				mockRFC8628CoreStorage.
+					EXPECT().
+					CreateUserCodeSession(ctx, "signature", gomock.Any()).
+					Return(nil).
+					Times(1)
+			},
+			check: func(t *testing.T, resp *fosite.DeviceResponse) {
+				assert.Equal(t, "userCode", resp.GetUserCode())
+			},
+		},
+		{
+			description: "should pass when generating a unique user code within allowed attempts",
+			setup: func() {
+				mockRFC8628CodeStrategy.
+					EXPECT().
+					GenerateDeviceCode(ctx).
+					Return("deviceCode", "signature", nil)
+				mockRFC8628CoreStorage.
+					EXPECT().
+					CreateDeviceCodeSession(ctx, "signature", gomock.Any()).
+					Return(nil)
+				gomock.InOrder(
+					mockRFC8628CodeStrategy.
+						EXPECT().
+						GenerateUserCode(ctx).
+						Return("duplicatedUserCode", "duplicatedSignature", nil),
+					mockRFC8628CoreStorage.
+						EXPECT().
+						CreateUserCodeSession(ctx, "duplicatedSignature", gomock.Any()).
+						Return(errors.New("unique constraint violation")),
+					mockRFC8628CodeStrategy.
+						EXPECT().
+						GenerateUserCode(ctx).
+						Return("uniqueUserCode", "uniqueSignature", nil),
+					mockRFC8628CoreStorage.
+						EXPECT().
+						CreateUserCodeSession(ctx, "uniqueSignature", gomock.Any()).
+						Return(nil),
+				)
+			},
+			check: func(t *testing.T, resp *fosite.DeviceResponse) {
+				assert.Equal(t, "uniqueUserCode", resp.GetUserCode())
+			},
+		},
+		{
+			description: "should fail after maximum retries to generate a unique user code",
+			setup: func() {
+				mockRFC8628CodeStrategy.
+					EXPECT().
+					GenerateDeviceCode(ctx).
+					Return("deviceCode", "signature", nil)
+				mockRFC8628CoreStorage.
+					EXPECT().
+					CreateDeviceCodeSession(ctx, "signature", gomock.Any()).
+					Return(nil)
+				mockRFC8628CodeStrategy.
+					EXPECT().
+					GenerateUserCode(ctx).
+					Return("duplicatedUserCode", "duplicatedSignature", nil).
+					Times(rfc8628.MaxAttempts)
+				mockRFC8628CoreStorage.
+					EXPECT().
+					CreateUserCodeSession(ctx, "duplicatedSignature", gomock.Any()).
+					Return(errors.New("unique constraint violation")).
+					Times(rfc8628.MaxAttempts)
+			},
+			check: func(t *testing.T, resp *fosite.DeviceResponse) {
+				assert.Empty(t, resp.GetUserCode())
+			},
+			expectError: fosite.ErrServerError,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(fmt.Sprintf("scenario=%s", testCase.description), func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			mockRFC8628CoreStorage = internal.NewMockRFC8628CoreStorage(ctrl)
+			mockRFC8628CodeStrategy = internal.NewMockRFC8628CodeStrategy(ctrl)
+
+			h := rfc8628.DeviceAuthHandler{
+				Storage:  mockRFC8628CoreStorage,
+				Strategy: mockRFC8628CodeStrategy,
+				Config: &fosite.Config{
+					DeviceAndUserCodeLifespan:      time.Minute * 10,
+					DeviceAuthTokenPollingInterval: time.Second * 5,
+					DeviceVerificationURL:          "www.test.com",
+					AccessTokenLifespan:            time.Hour,
+					RefreshTokenLifespan:           time.Hour,
+					ScopeStrategy:                  fosite.HierarchicScopeStrategy,
+					AudienceMatchingStrategy:       fosite.DefaultAudienceMatchingStrategy,
+					RefreshTokenScopes:             []string{"offline"},
+				},
+			}
+
+			if testCase.setup != nil {
+				testCase.setup()
+			}
+
+			resp := fosite.NewDeviceResponse()
+			err := h.HandleDeviceEndpointRequest(ctx, req, resp)
+
+			if testCase.expectError != nil {
+				require.EqualError(t, err, testCase.expectError.Error(), "%+v", err)
+			} else {
+				require.NoError(t, err, "%+v", err)
+			}
+
+			if testCase.check != nil {
+				testCase.check(t, resp)
+			}
+		})
+	}
+}