From 78bc6f2c1cdf6ce2c77015b2a5139cfb4ef9e742 Mon Sep 17 00:00:00 2001
From: Nicolas MASSE <nicolas.masse@itix.fr>
Date: Wed, 3 Feb 2021 09:15:11 +0100
Subject: [PATCH] fix race condition where realm import fails during token
 renew

---
 async/dispatcher.go | 52 ++++++++++++++++++++++++++++++++-------------
 async/renewer.go    | 20 +++++++++++++----
 2 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/async/dispatcher.go b/async/dispatcher.go
index a309f50..7349145 100644
--- a/async/dispatcher.go
+++ b/async/dispatcher.go
@@ -129,24 +129,31 @@ type Dispatcher struct {
 	users        chan KeycloakUserCreationRequest
 	Results      chan KeycloakResult
 	tokenRenewer TokenRenewer
+	expiredToken chan struct{}
+	newToken     chan string
 }
 
 func NewDispatcher(workers int, config keycloak.Config, credentials kcimport.KeycloakCredentials) (Dispatcher, error) {
 	var dispatcher Dispatcher
+	var err error
 
-	importer, err := kcimport.NewKeycloakImporter(config)
+	dispatcher.tokenRenewer, err = NewTokenRenewer(config)
 	if err != nil {
-		return dispatcher, err
+		return Dispatcher{}, err
 	}
-
-	importer.Credentials = credentials
-	err = importer.Login()
+	dispatcher.tokenRenewer.Importer.Credentials = credentials
+	err = dispatcher.tokenRenewer.Importer.Login()
 	if err != nil {
-		return dispatcher, err
+		return Dispatcher{}, err
 	}
 
-	dispatcher.tokenRenewer = NewTokenRenewer()
-	dispatcher.Importer = importer
+	dispatcher.Importer, err = kcimport.NewKeycloakImporter(config)
+	if err != nil {
+		return Dispatcher{}, err
+	}
+	dispatcher.Importer.Token = dispatcher.tokenRenewer.Importer.Token
+	dispatcher.expiredToken = dispatcher.tokenRenewer.expiredToken
+	dispatcher.newToken = make(chan string, 1)
 	dispatcher.clients = make(chan KeycloakClientCreationRequest)
 	dispatcher.users = make(chan KeycloakUserCreationRequest)
 	dispatcher.Results = make(chan KeycloakResult)
@@ -155,11 +162,11 @@ func NewDispatcher(workers int, config keycloak.Config, credentials kcimport.Key
 	for i := 0; i < workers; i++ {
 		dispatcher.Workers[i] = NewWorker(fmt.Sprintf("worker-%03d", i), dispatcher.clients, dispatcher.users, dispatcher.Results, dispatcher.tokenRenewer.expiredToken)
 
-		importer, err = kcimport.NewKeycloakImporter(config)
+		importer, err := kcimport.NewKeycloakImporter(config)
 		if err != nil {
-			return dispatcher, err
+			return Dispatcher{}, err
 		}
-		importer.Token = dispatcher.Importer.Token
+		importer.Token = dispatcher.tokenRenewer.Importer.Token
 
 		dispatcher.Workers[i].Importer = importer
 	}
@@ -169,17 +176,32 @@ func NewDispatcher(workers int, config keycloak.Config, credentials kcimport.Key
 
 func (dispatcher *Dispatcher) ApplyRealm(realm keycloak.RealmRepresentation) {
 	var err error
-
-	for i := 0; i < 3; i++ {
+	var retries int
+	for retries = 0; retries < 3; retries++ {
 		err = dispatcher.Importer.ApplyRealm(realm)
-		if err != nil {
-			continue
+		if err == nil {
+			break
+		}
+
+		if e, ok := err.(*kcimport.ImportError); ok {
+			if e.StatusCode == 401 {
+				dispatcher.expiredToken <- struct{}{}
+				select {
+				case newToken := <-dispatcher.newToken:
+					dispatcher.Importer.Token = newToken
+					continue
+				}
+			}
 		}
 	}
 
 	dispatcher.Results <- NewKeycloakResult("dispatcher", KeycloakRealm, realm.ID, nil, err, 0)
 }
 
+func (dispatcher *Dispatcher) NewToken(token string) {
+	dispatcher.newToken <- token
+}
+
 func (dispatcher *Dispatcher) ApplyClient(realmName string, client keycloak.ClientRepresentation) {
 	dispatcher.clients <- KeycloakClientCreationRequest{realmName, client}
 }
diff --git a/async/renewer.go b/async/renewer.go
index d14d7a4..5b43411 100644
--- a/async/renewer.go
+++ b/async/renewer.go
@@ -21,6 +21,9 @@ package async
 import (
 	"fmt"
 	"time"
+
+	keycloak "github.com/nmasse-itix/keycloak-client"
+	kcimport "github.com/nmasse-itix/keycloak-realm-import"
 )
 
 func (tr *TokenRenewer) RenewToken(dispatcher *Dispatcher) {
@@ -33,7 +36,7 @@ func (tr *TokenRenewer) RenewToken(dispatcher *Dispatcher) {
 				continue
 			}
 
-			err := dispatcher.Importer.Login()
+			err := tr.Importer.Login()
 			if err != nil {
 				fmt.Printf("dispatcher: Cannot renew OIDC token: %s\n", err)
 				continue
@@ -41,8 +44,9 @@ func (tr *TokenRenewer) RenewToken(dispatcher *Dispatcher) {
 
 			tr.LastTokenRenew = time.Now()
 			for i := 0; i < len(dispatcher.Workers); i++ {
-				dispatcher.Workers[i].NewToken(dispatcher.Importer.Token)
+				dispatcher.Workers[i].NewToken(tr.Importer.Token)
 			}
+			dispatcher.NewToken(tr.Importer.Token)
 		}
 	}
 }
@@ -54,13 +58,21 @@ func (tr *TokenRenewer) Stop() {
 type TokenRenewer struct {
 	quit           chan struct{}
 	expiredToken   chan struct{}
+	Importer       kcimport.KeycloakImporter
 	LastTokenRenew time.Time
 }
 
-func NewTokenRenewer() TokenRenewer {
+func NewTokenRenewer(config keycloak.Config) (TokenRenewer, error) {
 	var tokenRenewer TokenRenewer
 	tokenRenewer.LastTokenRenew = time.Now()
 	tokenRenewer.expiredToken = make(chan struct{})
 	tokenRenewer.quit = make(chan struct{})
-	return tokenRenewer
+
+	var err error
+	tokenRenewer.Importer, err = kcimport.NewKeycloakImporter(config)
+	if err != nil {
+		return TokenRenewer{}, err
+	}
+
+	return tokenRenewer, nil
 }