fix: move owner mutex to userConfig

internal/broker/broker.go

@@ -49,8 +49,6 @@ type Config struct {
 type Broker struct {
 	cfg Config
 
-	setOwnerMutex sync.Mutex
-
 	provider providers.Provider
 	oidcCfg  oidc.Config
 
@@ -615,8 +613,7 @@ func (b *Broker) handleIsAuthenticated(ctx context.Context, session *session, au
 		}
 	}
 
-	err = b.maybeRegisterOwner(b.cfg.ConfigFile, authInfo.UserInfo.Name)
-	if err != nil {
+	if err := b.cfg.registerOwner(b.cfg.ConfigFile, authInfo.UserInfo.Name); err != nil {
 		// The user is not allowed if we fail to create the owner-autoregistration file.
 		// Otherwise the owner might change if the broker is restarted.
 		slog.Error(fmt.Sprintf("Failed to assign the owner role: %v", err))
@@ -643,19 +640,6 @@ func (b *Broker) handleIsAuthenticated(ctx context.Context, session *session, au
 	return AuthGranted, userInfoMessage{UserInfo: authInfo.UserInfo}
 }
 
-func (b *Broker) maybeRegisterOwner(cfgPath string, userName string) error {
-	// We need to lock here to avoid a race condition where two users log in at the same time, causing both to be
-	// considered the owner.
-	b.setOwnerMutex.Lock()
-	defer b.setOwnerMutex.Unlock()
-
-	if !b.cfg.userConfig.shouldRegisterOwner() {
-		return nil
-	}
-
-	return b.cfg.registerOwner(cfgPath, userName)
-}
-
 // userNameIsAllowed checks whether the user's username is allowed to access the machine.
 func (b *Broker) userNameIsAllowed(userName string) bool {
 	normalizedUsername := b.provider.NormalizeUsername(userName)
@@ -669,7 +653,8 @@ func (b *Broker) userNameIsAllowed(userName string) bool {
 	if _, ok := b.cfg.userConfig.allowedUsers[normalizedUsername]; ok {
 		return true
 	}
-	return b.cfg.userConfig.ownerAllowed && b.cfg.userConfig.owner == normalizedUsername
+
+	return b.cfg.isOwnerAllowed(normalizedUsername)
 }
 
 func (b *Broker) startAuthenticate(sessionID string) (context.Context, error) {

internal/broker/config.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"gopkg.in/ini.v1"
 )
@@ -68,6 +69,7 @@ type userConfig struct {
 	ownerAllowed          bool
 	firstUserBecomesOwner bool
 	owner                 string
+	ownerMutex            *sync.RWMutex
 	homeBaseDir           string
 	allowedSSHSuffixes    []string
 
@@ -103,6 +105,9 @@ func getDropInFiles(cfgPath string) ([]any, error) {
 }
 
 func (uc *userConfig) populateUsersConfig(users *ini.Section) {
+	uc.ownerMutex.Lock()
+	defer uc.ownerMutex.Unlock()
+
 	if users == nil {
 		// The default behavior is to allow only the owner
 		uc.ownerAllowed = true
@@ -146,7 +151,7 @@ func (uc *userConfig) populateUsersConfig(users *ini.Section) {
 
 // parseConfigFile parses the config file and returns a map with the configuration keys and values.
 func parseConfigFile(cfgPath string, p provider) (userConfig, error) {
-	cfg := userConfig{provider: p}
+	cfg := userConfig{provider: p, ownerMutex: &sync.RWMutex{}}
 
 	dropInFiles, err := getDropInFiles(cfgPath)
 	if err != nil {
@@ -182,11 +187,23 @@ func parseConfigFile(cfgPath string, p provider) (userConfig, error) {
 	return cfg, nil
 }
 
-func (uc *userConfig) shouldRegisterOwner() bool {
-	return uc.ownerAllowed && uc.firstUserBecomesOwner && uc.owner == ""
+func (uc *userConfig) isOwnerAllowed(userName string) bool {
+	uc.ownerMutex.RLock()
+	defer uc.ownerMutex.RUnlock()
+
+	return uc.ownerAllowed && uc.owner == userName
 }
 
 func (uc *userConfig) registerOwner(cfgPath, userName string) error {
+	// We need to lock here to avoid a race condition where two users log in at the same time, causing both to be
+	// considered the owner.
+	uc.ownerMutex.Lock()
+	defer uc.ownerMutex.Unlock()
+
+	if shouldRegister := uc.ownerAllowed && uc.firstUserBecomesOwner && uc.owner == ""; !shouldRegister {
+		return nil
+	}
+
 	if cfgPath == "" {
 		uc.owner = uc.provider.NormalizeUsername(userName)
 		uc.firstUserBecomesOwner = false

internal/broker/config_test.go

@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"strings"
+	"sync"
 	"testing"
 	"unsafe"
 
@@ -58,7 +59,7 @@ issuer = https://higher-precedence-issuer.url.com
 func TestParseConfig(t *testing.T) {
 	t.Parallel()
 	p := &testutils.MockProvider{}
-	uncheckedFields := map[string]struct{}{"provider": {}}
+	ignoredFields := map[string]struct{}{"provider": {}, "ownerMutex": {}}
 
 	tests := map[string]struct {
 		configType string
@@ -140,7 +141,7 @@ func TestParseConfig(t *testing.T) {
 			typ := reflect.TypeOf(&cfg).Elem()
 			for i := 0; i < typ.NumField(); i++ {
 				field := typ.Field(i)
-				if _, ok := uncheckedFields[field.Name]; ok {
+				if _, ok := ignoredFields[field.Name]; ok {
 					continue
 				}
 				fieldValue := val.Field(i)
@@ -318,7 +319,7 @@ func TestRegisterOwner(t *testing.T) {
 	err = os.Mkdir(dropInDir, 0700)
 	require.NoError(t, err, "Setup: Failed to create drop-in directory")
 
-	cfg := userConfig{firstUserBecomesOwner: true, ownerAllowed: true, provider: p}
+	cfg := userConfig{firstUserBecomesOwner: true, ownerAllowed: true, provider: p, ownerMutex: &sync.RWMutex{}}
 	err = cfg.registerOwner(confPath, userName)
 	require.NoError(t, err)
 

internal/broker/export_test.go

@@ -2,11 +2,16 @@ package broker
 
 import (
 	"context"
+	"sync"
 
 	"github.com/ubuntu/authd-oidc-brokers/internal/providers/info"
 	tokenPkg "github.com/ubuntu/authd-oidc-brokers/internal/token"
 )
 
+func (cfg *Config) Init() {
+	cfg.ownerMutex = &sync.RWMutex{}
+}
+
 func (cfg *Config) SetClientID(clientID string) {
 	cfg.clientID = clientID
 }
@@ -24,10 +29,16 @@ func (cfg *Config) SetAllowedUsers(allowedUsers map[string]struct{}) {
 }
 
 func (cfg *Config) SetOwner(owner string) {
+	cfg.ownerMutex.Lock()
+	defer cfg.ownerMutex.Unlock()
+
 	cfg.owner = owner
 }
 
 func (cfg *Config) SetFirstUserBecomesOwner(firstUserBecomesOwner bool) {
+	cfg.ownerMutex.Lock()
+	defer cfg.ownerMutex.Unlock()
+
 	cfg.firstUserBecomesOwner = firstUserBecomesOwner
 }
 
@@ -36,6 +47,9 @@ func (cfg *Config) SetAllUsersAllowed(allUsersAllowed bool) {
 }
 
 func (cfg *Config) SetOwnerAllowed(ownerAllowed bool) {
+	cfg.ownerMutex.Lock()
+	defer cfg.ownerMutex.Unlock()
+
 	cfg.ownerAllowed = ownerAllowed
 }
 

internal/broker/helper_test.go

@@ -47,6 +47,7 @@ type brokerForTestConfig struct {
 func newBrokerForTests(t *testing.T, cfg *brokerForTestConfig) (b *broker.Broker) {
 	t.Helper()
 
+	cfg.Init()
 	if cfg.issuerURL != "" {
 		cfg.SetIssuerURL(cfg.issuerURL)
 	}