Add context to throughout the project and update logging (#2209)

propagate context from all the API calls and log request ID, account ID and peer ID

---------

Co-authored-by: Zoltan Papp <zoltan.pmail@gmail.com>

management/server/idp/auth0.go

@@ -183,7 +183,7 @@ func (c *Auth0Credentials) jwtStillValid() bool {
 }
 
 // requestJWTToken performs request to get jwt token
-func (c *Auth0Credentials) requestJWTToken() (*http.Response, error) {
+func (c *Auth0Credentials) requestJWTToken(ctx context.Context) (*http.Response, error) {
 	var res *http.Response
 	reqURL := c.clientConfig.AuthIssuer + "/oauth/token"
 
@@ -200,7 +200,7 @@ func (c *Auth0Credentials) requestJWTToken() (*http.Response, error) {
 
 	req.Header.Add("content-type", "application/json")
 
-	log.Debug("requesting new jwt token for idp manager")
+	log.WithContext(ctx).Debug("requesting new jwt token for idp manager")
 
 	res, err = c.httpClient.Do(req)
 	if err != nil {
@@ -247,7 +247,7 @@ func (c *Auth0Credentials) parseRequestJWTResponse(rawBody io.ReadCloser) (JWTTo
 }
 
 // Authenticate retrieves access token to use the Auth0 Management API
-func (c *Auth0Credentials) Authenticate() (JWTToken, error) {
+func (c *Auth0Credentials) Authenticate(ctx context.Context) (JWTToken, error) {
 	c.mux.Lock()
 	defer c.mux.Unlock()
 
@@ -260,14 +260,14 @@ func (c *Auth0Credentials) Authenticate() (JWTToken, error) {
 		return c.jwtToken, nil
 	}
 
-	res, err := c.requestJWTToken()
+	res, err := c.requestJWTToken(ctx)
 	if err != nil {
 		return c.jwtToken, err
 	}
 	defer func() {
 		err = res.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing get jwt token response body: %v", err)
+			log.WithContext(ctx).Errorf("error while closing get jwt token response body: %v", err)
 		}
 	}()
 
@@ -301,8 +301,8 @@ func requestByUserIDURL(authIssuer, userID string) string {
 }
 
 // GetAccount returns all the users for a given profile. Calls Auth0 API.
-func (am *Auth0Manager) GetAccount(accountID string) ([]*UserData, error) {
-	jwtToken, err := am.credentials.Authenticate()
+func (am *Auth0Manager) GetAccount(ctx context.Context, accountID string) ([]*UserData, error) {
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -353,7 +353,7 @@ func (am *Auth0Manager) GetAccount(accountID string) ([]*UserData, error) {
 			return nil, err
 		}
 
-		log.Debugf("returned user batch for accountID %s on page %d, batch length %d", accountID, page, len(batch))
+		log.WithContext(ctx).Debugf("returned user batch for accountID %s on page %d, batch length %d", accountID, page, len(batch))
 
 		err = res.Body.Close()
 		if err != nil {
@@ -365,7 +365,7 @@ func (am *Auth0Manager) GetAccount(accountID string) ([]*UserData, error) {
 		}
 
 		if len(batch) == 0 || len(batch) < resultsPerPage {
-			log.Debugf("finished loading users for accountID %s", accountID)
+			log.WithContext(ctx).Debugf("finished loading users for accountID %s", accountID)
 			return list, nil
 		}
 	}
@@ -374,8 +374,8 @@ func (am *Auth0Manager) GetAccount(accountID string) ([]*UserData, error) {
 }
 
 // GetUserDataByID requests user data from auth0 via ID
-func (am *Auth0Manager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
-	jwtToken, err := am.credentials.Authenticate()
+func (am *Auth0Manager) GetUserDataByID(ctx context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -414,7 +414,7 @@ func (am *Auth0Manager) GetUserDataByID(userID string, appMetadata AppMetadata)
 	defer func() {
 		err = res.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing update user app metadata response body: %v", err)
+			log.WithContext(ctx).Errorf("error while closing update user app metadata response body: %v", err)
 		}
 	}()
 
@@ -426,9 +426,9 @@ func (am *Auth0Manager) GetUserDataByID(userID string, appMetadata AppMetadata)
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userId and metadata map
-func (am *Auth0Manager) UpdateUserAppMetadata(userID string, appMetadata AppMetadata) error {
+func (am *Auth0Manager) UpdateUserAppMetadata(ctx context.Context, userID string, appMetadata AppMetadata) error {
 
-	jwtToken, err := am.credentials.Authenticate()
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return err
 	}
@@ -449,7 +449,7 @@ func (am *Auth0Manager) UpdateUserAppMetadata(userID string, appMetadata AppMeta
 	req.Header.Add("authorization", "Bearer "+jwtToken.AccessToken)
 	req.Header.Add("content-type", "application/json")
 
-	log.Debugf("updating IdP metadata for user %s", userID)
+	log.WithContext(ctx).Debugf("updating IdP metadata for user %s", userID)
 
 	res, err := am.httpClient.Do(req)
 	if err != nil {
@@ -466,7 +466,7 @@ func (am *Auth0Manager) UpdateUserAppMetadata(userID string, appMetadata AppMeta
 	defer func() {
 		err = res.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing update user app metadata response body: %v", err)
+			log.WithContext(ctx).Errorf("error while closing update user app metadata response body: %v", err)
 		}
 	}()
 
@@ -530,9 +530,9 @@ func buildUserExportRequest() (string, error) {
 }
 
 func (am *Auth0Manager) createRequest(
-	method string, endpoint string, body io.Reader,
+	ctx context.Context, method string, endpoint string, body io.Reader,
 ) (*http.Request, error) {
-	jwtToken, err := am.credentials.Authenticate()
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -548,8 +548,8 @@ func (am *Auth0Manager) createRequest(
 	return req, nil
 }
 
-func (am *Auth0Manager) createPostRequest(endpoint string, payloadStr string) (*http.Request, error) {
-	req, err := am.createRequest("POST", endpoint, strings.NewReader(payloadStr))
+func (am *Auth0Manager) createPostRequest(ctx context.Context, endpoint string, payloadStr string) (*http.Request, error) {
+	req, err := am.createRequest(ctx, "POST", endpoint, strings.NewReader(payloadStr))
 	if err != nil {
 		return nil, err
 	}
@@ -560,20 +560,20 @@ func (am *Auth0Manager) createPostRequest(endpoint string, payloadStr string) (*
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (am *Auth0Manager) GetAllAccounts() (map[string][]*UserData, error) {
+func (am *Auth0Manager) GetAllAccounts(ctx context.Context) (map[string][]*UserData, error) {
 	payloadString, err := buildUserExportRequest()
 	if err != nil {
 		return nil, err
 	}
 
-	exportJobReq, err := am.createPostRequest("/api/v2/jobs/users-exports", payloadString)
+	exportJobReq, err := am.createPostRequest(ctx, "/api/v2/jobs/users-exports", payloadString)
 	if err != nil {
 		return nil, err
 	}
 
 	jobResp, err := am.httpClient.Do(exportJobReq)
 	if err != nil {
-		log.Debugf("Couldn't get job response %v", err)
+		log.WithContext(ctx).Debugf("Couldn't get job response %v", err)
 		if am.appMetrics != nil {
 			am.appMetrics.IDPMetrics().CountRequestError()
 		}
@@ -583,7 +583,7 @@ func (am *Auth0Manager) GetAllAccounts() (map[string][]*UserData, error) {
 	defer func() {
 		err = jobResp.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing update user app metadata response body: %v", err)
+			log.WithContext(ctx).Errorf("error while closing update user app metadata response body: %v", err)
 		}
 	}()
 	if jobResp.StatusCode != 200 {
@@ -597,13 +597,13 @@ func (am *Auth0Manager) GetAllAccounts() (map[string][]*UserData, error) {
 
 	body, err := io.ReadAll(jobResp.Body)
 	if err != nil {
-		log.Debugf("Couldn't read export job response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't read export job response; %v", err)
 		return nil, err
 	}
 
 	err = am.helper.Unmarshal(body, &exportJobResp)
 	if err != nil {
-		log.Debugf("Couldn't unmarshal export job response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't unmarshal export job response; %v", err)
 		return nil, err
 	}
 
@@ -614,16 +614,16 @@ func (am *Auth0Manager) GetAllAccounts() (map[string][]*UserData, error) {
 		return nil, fmt.Errorf("couldn't get an batch id status %d, %s, response body: %v", jobResp.StatusCode, jobResp.Status, exportJobResp)
 	}
 
-	log.Debugf("batch id status %d, %s, response body: %v", jobResp.StatusCode, jobResp.Status, exportJobResp)
+	log.WithContext(ctx).Debugf("batch id status %d, %s, response body: %v", jobResp.StatusCode, jobResp.Status, exportJobResp)
 
-	done, downloadLink, err := am.checkExportJobStatus(exportJobResp.ID)
+	done, downloadLink, err := am.checkExportJobStatus(ctx, exportJobResp.ID)
 	if err != nil {
-		log.Debugf("Failed at getting status checks from exportJob; %v", err)
+		log.WithContext(ctx).Debugf("Failed at getting status checks from exportJob; %v", err)
 		return nil, err
 	}
 
 	if done {
-		return am.downloadProfileExport(downloadLink)
+		return am.downloadProfileExport(ctx, downloadLink)
 	}
 
 	return nil, fmt.Errorf("failed extracting user profiles from auth0")
@@ -632,13 +632,13 @@ func (am *Auth0Manager) GetAllAccounts() (map[string][]*UserData, error) {
 // GetUserByEmail searches users with a given email. If no users have been found, this function returns an empty list.
 // This function can return multiple users. This is due to the Auth0 internals - there could be multiple users with
 // the same email but different connections that are considered as separate accounts (e.g., Google and username/password).
-func (am *Auth0Manager) GetUserByEmail(email string) ([]*UserData, error) {
-	jwtToken, err := am.credentials.Authenticate()
+func (am *Auth0Manager) GetUserByEmail(ctx context.Context, email string) ([]*UserData, error) {
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
 	reqURL := am.authIssuer + "/api/v2/users-by-email?email=" + url.QueryEscape(email)
-	body, err := doGetReq(am.httpClient, reqURL, jwtToken.AccessToken)
+	body, err := doGetReq(ctx, am.httpClient, reqURL, jwtToken.AccessToken)
 	if err != nil {
 		return nil, err
 	}
@@ -651,7 +651,7 @@ func (am *Auth0Manager) GetUserByEmail(email string) ([]*UserData, error) {
 
 	err = am.helper.Unmarshal(body, &userResp)
 	if err != nil {
-		log.Debugf("Couldn't unmarshal export job response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't unmarshal export job response; %v", err)
 		return nil, err
 	}
 
@@ -659,13 +659,13 @@ func (am *Auth0Manager) GetUserByEmail(email string) ([]*UserData, error) {
 }
 
 // CreateUser creates a new user in Auth0 Idp and sends an invite
-func (am *Auth0Manager) CreateUser(email, name, accountID, invitedByEmail string) (*UserData, error) {
+func (am *Auth0Manager) CreateUser(ctx context.Context, email, name, accountID, invitedByEmail string) (*UserData, error) {
 
 	payloadString, err := buildCreateUserRequestPayload(email, name, accountID, invitedByEmail)
 	if err != nil {
 		return nil, err
 	}
-	req, err := am.createPostRequest("/api/v2/users", payloadString)
+	req, err := am.createPostRequest(ctx, "/api/v2/users", payloadString)
 	if err != nil {
 		return nil, err
 	}
@@ -676,7 +676,7 @@ func (am *Auth0Manager) CreateUser(email, name, accountID, invitedByEmail string
 
 	resp, err := am.httpClient.Do(req)
 	if err != nil {
-		log.Debugf("Couldn't get job response %v", err)
+		log.WithContext(ctx).Debugf("Couldn't get job response %v", err)
 		if am.appMetrics != nil {
 			am.appMetrics.IDPMetrics().CountRequestError()
 		}
@@ -686,7 +686,7 @@ func (am *Auth0Manager) CreateUser(email, name, accountID, invitedByEmail string
 	defer func() {
 		err = resp.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing create user response body: %v", err)
+			log.WithContext(ctx).Errorf("error while closing create user response body: %v", err)
 		}
 	}()
 	if !(resp.StatusCode == 200 || resp.StatusCode == 201) {
@@ -700,13 +700,13 @@ func (am *Auth0Manager) CreateUser(email, name, accountID, invitedByEmail string
 
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
-		log.Debugf("Couldn't read export job response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't read export job response; %v", err)
 		return nil, err
 	}
 
 	err = am.helper.Unmarshal(body, &createResp)
 	if err != nil {
-		log.Debugf("Couldn't unmarshal export job response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't unmarshal export job response; %v", err)
 		return nil, err
 	}
 
@@ -714,14 +714,14 @@ func (am *Auth0Manager) CreateUser(email, name, accountID, invitedByEmail string
 		return nil, fmt.Errorf("couldn't create user: response %v", resp)
 	}
 
-	log.Debugf("created user %s in account %s", createResp.ID, accountID)
+	log.WithContext(ctx).Debugf("created user %s in account %s", createResp.ID, accountID)
 
 	return &createResp, nil
 }
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (am *Auth0Manager) InviteUserByID(userID string) error {
+func (am *Auth0Manager) InviteUserByID(ctx context.Context, userID string) error {
 	userVerificationReq := userVerificationJobRequest{
 		UserID: userID,
 	}
@@ -731,14 +731,14 @@ func (am *Auth0Manager) InviteUserByID(userID string) error {
 		return err
 	}
 
-	req, err := am.createPostRequest("/api/v2/jobs/verification-email", string(payload))
+	req, err := am.createPostRequest(ctx, "/api/v2/jobs/verification-email", string(payload))
 	if err != nil {
 		return err
 	}
 
 	resp, err := am.httpClient.Do(req)
 	if err != nil {
-		log.Debugf("Couldn't get job response %v", err)
+		log.WithContext(ctx).Debugf("Couldn't get job response %v", err)
 		if am.appMetrics != nil {
 			am.appMetrics.IDPMetrics().CountRequestError()
 		}
@@ -748,7 +748,7 @@ func (am *Auth0Manager) InviteUserByID(userID string) error {
 	defer func() {
 		err = resp.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing invite user response body: %v", err)
+			log.WithContext(ctx).Errorf("error while closing invite user response body: %v", err)
 		}
 	}()
 	if !(resp.StatusCode == 200 || resp.StatusCode == 201) {
@@ -762,15 +762,15 @@ func (am *Auth0Manager) InviteUserByID(userID string) error {
 }
 
 // DeleteUser from Auth0
-func (am *Auth0Manager) DeleteUser(userID string) error {
-	req, err := am.createRequest(http.MethodDelete, "/api/v2/users/"+url.QueryEscape(userID), nil)
+func (am *Auth0Manager) DeleteUser(ctx context.Context, userID string) error {
+	req, err := am.createRequest(ctx, http.MethodDelete, "/api/v2/users/"+url.QueryEscape(userID), nil)
 	if err != nil {
 		return err
 	}
 
 	resp, err := am.httpClient.Do(req)
 	if err != nil {
-		log.Debugf("execute delete request: %v", err)
+		log.WithContext(ctx).Debugf("execute delete request: %v", err)
 		if am.appMetrics != nil {
 			am.appMetrics.IDPMetrics().CountRequestError()
 		}
@@ -780,7 +780,7 @@ func (am *Auth0Manager) DeleteUser(userID string) error {
 	defer func() {
 		err = resp.Body.Close()
 		if err != nil {
-			log.Errorf("close delete request body: %v", err)
+			log.WithContext(ctx).Errorf("close delete request body: %v", err)
 		}
 	}()
 	if resp.StatusCode != 204 {
@@ -795,20 +795,20 @@ func (am *Auth0Manager) DeleteUser(userID string) error {
 
 // GetAllConnections returns detailed list of all connections filtered by given params.
 // Note this method is not part of the IDP Manager interface as this is Auth0 specific.
-func (am *Auth0Manager) GetAllConnections(strategy []string) ([]Connection, error) {
+func (am *Auth0Manager) GetAllConnections(ctx context.Context, strategy []string) ([]Connection, error) {
 	var connections []Connection
 
 	q := make(url.Values)
 	q.Set("strategy", strings.Join(strategy, ","))
 
-	req, err := am.createRequest(http.MethodGet, "/api/v2/connections?"+q.Encode(), nil)
+	req, err := am.createRequest(ctx, http.MethodGet, "/api/v2/connections?"+q.Encode(), nil)
 	if err != nil {
 		return connections, err
 	}
 
 	resp, err := am.httpClient.Do(req)
 	if err != nil {
-		log.Debugf("execute get connections request: %v", err)
+		log.WithContext(ctx).Debugf("execute get connections request: %v", err)
 		if am.appMetrics != nil {
 			am.appMetrics.IDPMetrics().CountRequestError()
 		}
@@ -818,7 +818,7 @@ func (am *Auth0Manager) GetAllConnections(strategy []string) ([]Connection, erro
 	defer func() {
 		err = resp.Body.Close()
 		if err != nil {
-			log.Errorf("close get connections request body: %v", err)
+			log.WithContext(ctx).Errorf("close get connections request body: %v", err)
 		}
 	}()
 	if resp.StatusCode != 200 {
@@ -830,13 +830,13 @@ func (am *Auth0Manager) GetAllConnections(strategy []string) ([]Connection, erro
 
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
-		log.Debugf("Couldn't read get connections response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't read get connections response; %v", err)
 		return connections, err
 	}
 
 	err = am.helper.Unmarshal(body, &connections)
 	if err != nil {
-		log.Debugf("Couldn't unmarshal get connection response; %v", err)
+		log.WithContext(ctx).Debugf("Couldn't unmarshal get connection response; %v", err)
 		return connections, err
 	}
 
@@ -845,23 +845,23 @@ func (am *Auth0Manager) GetAllConnections(strategy []string) ([]Connection, erro
 
 // checkExportJobStatus checks the status of the job created at CreateExportUsersJob.
 // If the status is "completed", then return the downloadLink
-func (am *Auth0Manager) checkExportJobStatus(jobID string) (bool, string, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+func (am *Auth0Manager) checkExportJobStatus(ctx context.Context, jobID string) (bool, string, error) {
+	ctx, cancel := context.WithTimeout(ctx, 90*time.Second)
 	defer cancel()
 	retry := time.NewTicker(10 * time.Second)
 	for {
 		select {
 		case <-ctx.Done():
-			log.Debugf("Export job status stopped...\n")
+			log.WithContext(ctx).Debugf("Export job status stopped...\n")
 			return false, "", ctx.Err()
 		case <-retry.C:
-			jwtToken, err := am.credentials.Authenticate()
+			jwtToken, err := am.credentials.Authenticate(ctx)
 			if err != nil {
 				return false, "", err
 			}
 
 			statusURL := am.authIssuer + "/api/v2/jobs/" + jobID
-			body, err := doGetReq(am.httpClient, statusURL, jwtToken.AccessToken)
+			body, err := doGetReq(ctx, am.httpClient, statusURL, jwtToken.AccessToken)
 			if err != nil {
 				return false, "", err
 			}
@@ -872,7 +872,7 @@ func (am *Auth0Manager) checkExportJobStatus(jobID string) (bool, string, error)
 				return false, "", err
 			}
 
-			log.Debugf("current export job status is %v", status.Status)
+			log.WithContext(ctx).Debugf("current export job status is %v", status.Status)
 
 			if status.Status != "completed" {
 				continue
@@ -884,8 +884,8 @@ func (am *Auth0Manager) checkExportJobStatus(jobID string) (bool, string, error)
 }
 
 // downloadProfileExport downloads user profiles from auth0 batch job
-func (am *Auth0Manager) downloadProfileExport(location string) (map[string][]*UserData, error) {
-	body, err := doGetReq(am.httpClient, location, "")
+func (am *Auth0Manager) downloadProfileExport(ctx context.Context, location string) (map[string][]*UserData, error) {
+	body, err := doGetReq(ctx, am.httpClient, location, "")
 	if err != nil {
 		return nil, err
 	}
@@ -927,7 +927,7 @@ func (am *Auth0Manager) downloadProfileExport(location string) (map[string][]*Us
 }
 
 // Boilerplate implementation for Get Requests.
-func doGetReq(client ManagerHTTPClient, url, accessToken string) ([]byte, error) {
+func doGetReq(ctx context.Context, client ManagerHTTPClient, url, accessToken string) ([]byte, error) {
 	req, err := http.NewRequest("GET", url, nil)
 	if err != nil {
 		return nil, err
@@ -945,7 +945,7 @@ func doGetReq(client ManagerHTTPClient, url, accessToken string) ([]byte, error)
 	defer func() {
 		err = res.Body.Close()
 		if err != nil {
-			log.Errorf("error while closing body for url %s: %v", url, err)
+			log.WithContext(ctx).Errorf("error while closing body for url %s: %v", url, err)
 		}
 	}()
 	body, err := io.ReadAll(res.Body)

management/server/idp/auth0_test.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -60,7 +61,7 @@ type mockAuth0Credentials struct {
 	err      error
 }
 
-func (mc *mockAuth0Credentials) Authenticate() (JWTToken, error) {
+func (mc *mockAuth0Credentials) Authenticate(_ context.Context) (JWTToken, error) {
 	return mc.jwtToken, mc.err
 }
 
@@ -126,7 +127,7 @@ func TestAuth0_RequestJWTToken(t *testing.T) {
 				helper:       testCase.helper,
 			}
 
-			res, err := creds.requestJWTToken()
+			res, err := creds.requestJWTToken(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")
@@ -295,7 +296,7 @@ func TestAuth0_Authenticate(t *testing.T) {
 
 			creds.jwtToken.expiresInTime = testCase.inputExpireToken
 
-			_, err := creds.Authenticate()
+			_, err := creds.Authenticate(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")
@@ -417,7 +418,7 @@ func TestAuth0_UpdateUserAppMetadata(t *testing.T) {
 				helper:      testCase.helper,
 			}
 
-			err := manager.UpdateUserAppMetadata("1", testCase.appMetadata)
+			err := manager.UpdateUserAppMetadata(context.Background(), "1", testCase.appMetadata)
 			testCase.assertErrFunc(t, err, testCase.assertErrFuncMessage)
 
 			assert.Equal(t, testCase.expectedReqBody, jwtReqClient.reqBody, "request body should match")

management/server/idp/authentik.go

@@ -116,7 +116,7 @@ func (ac *AuthentikCredentials) jwtStillValid() bool {
 }
 
 // requestJWTToken performs request to get jwt token.
-func (ac *AuthentikCredentials) requestJWTToken() (*http.Response, error) {
+func (ac *AuthentikCredentials) requestJWTToken(ctx context.Context) (*http.Response, error) {
 	data := url.Values{}
 	data.Set("client_id", ac.clientConfig.ClientID)
 	data.Set("username", ac.clientConfig.Username)
@@ -131,7 +131,7 @@ func (ac *AuthentikCredentials) requestJWTToken() (*http.Response, error) {
 	}
 	req.Header.Add("content-type", "application/x-www-form-urlencoded")
 
-	log.Debug("requesting new jwt token for authentik idp manager")
+	log.WithContext(ctx).Debug("requesting new jwt token for authentik idp manager")
 
 	resp, err := ac.httpClient.Do(req)
 	if err != nil {
@@ -183,7 +183,7 @@ func (ac *AuthentikCredentials) parseRequestJWTResponse(rawBody io.ReadCloser) (
 }
 
 // Authenticate retrieves access token to use the authentik management API.
-func (ac *AuthentikCredentials) Authenticate() (JWTToken, error) {
+func (ac *AuthentikCredentials) Authenticate(ctx context.Context) (JWTToken, error) {
 	ac.mux.Lock()
 	defer ac.mux.Unlock()
 
@@ -197,7 +197,7 @@ func (ac *AuthentikCredentials) Authenticate() (JWTToken, error) {
 		return ac.jwtToken, nil
 	}
 
-	resp, err := ac.requestJWTToken()
+	resp, err := ac.requestJWTToken(ctx)
 	if err != nil {
 		return ac.jwtToken, err
 	}
@@ -214,13 +214,13 @@ func (ac *AuthentikCredentials) Authenticate() (JWTToken, error) {
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userID and metadata map.
-func (am *AuthentikManager) UpdateUserAppMetadata(_ string, _ AppMetadata) error {
+func (am *AuthentikManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
 // GetUserDataByID requests user data from authentik via ID.
-func (am *AuthentikManager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
-	ctx, err := am.authenticationContext()
+func (am *AuthentikManager) GetUserDataByID(ctx context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
+	ctx, err := am.authenticationContext(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -254,8 +254,8 @@ func (am *AuthentikManager) GetUserDataByID(userID string, appMetadata AppMetada
 }
 
 // GetAccount returns all the users for a given profile.
-func (am *AuthentikManager) GetAccount(accountID string) ([]*UserData, error) {
-	users, err := am.getAllUsers()
+func (am *AuthentikManager) GetAccount(ctx context.Context, accountID string) ([]*UserData, error) {
+	users, err := am.getAllUsers(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -274,8 +274,8 @@ func (am *AuthentikManager) GetAccount(accountID string) ([]*UserData, error) {
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (am *AuthentikManager) GetAllAccounts() (map[string][]*UserData, error) {
-	users, err := am.getAllUsers()
+func (am *AuthentikManager) GetAllAccounts(ctx context.Context) (map[string][]*UserData, error) {
+	users, err := am.getAllUsers(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -291,12 +291,12 @@ func (am *AuthentikManager) GetAllAccounts() (map[string][]*UserData, error) {
 }
 
 // getAllUsers returns all users in a Authentik account.
-func (am *AuthentikManager) getAllUsers() ([]*UserData, error) {
+func (am *AuthentikManager) getAllUsers(ctx context.Context) ([]*UserData, error) {
 	users := make([]*UserData, 0)
 
 	page := int32(1)
 	for {
-		ctx, err := am.authenticationContext()
+		ctx, err := am.authenticationContext(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -329,14 +329,14 @@ func (am *AuthentikManager) getAllUsers() ([]*UserData, error) {
 }
 
 // CreateUser creates a new user in authentik Idp and sends an invitation.
-func (am *AuthentikManager) CreateUser(_, _, _, _ string) (*UserData, error) {
+func (am *AuthentikManager) CreateUser(_ context.Context, _, _, _, _ string) (*UserData, error) {
 	return nil, fmt.Errorf("method CreateUser not implemented")
 }
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (am *AuthentikManager) GetUserByEmail(email string) ([]*UserData, error) {
-	ctx, err := am.authenticationContext()
+func (am *AuthentikManager) GetUserByEmail(ctx context.Context, email string) ([]*UserData, error) {
+	ctx, err := am.authenticationContext(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -368,13 +368,13 @@ func (am *AuthentikManager) GetUserByEmail(email string) ([]*UserData, error) {
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (am *AuthentikManager) InviteUserByID(_ string) error {
+func (am *AuthentikManager) InviteUserByID(_ context.Context, _ string) error {
 	return fmt.Errorf("method InviteUserByID not implemented")
 }
 
 // DeleteUser from Authentik
-func (am *AuthentikManager) DeleteUser(userID string) error {
-	ctx, err := am.authenticationContext()
+func (am *AuthentikManager) DeleteUser(ctx context.Context, userID string) error {
+	ctx, err := am.authenticationContext(ctx)
 	if err != nil {
 		return err
 	}
@@ -404,8 +404,8 @@ func (am *AuthentikManager) DeleteUser(userID string) error {
 	return nil
 }
 
-func (am *AuthentikManager) authenticationContext() (context.Context, error) {
-	jwtToken, err := am.credentials.Authenticate()
+func (am *AuthentikManager) authenticationContext(ctx context.Context) (context.Context, error) {
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}

management/server/idp/authentik_test.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"strings"
@@ -138,7 +139,7 @@ func TestAuthentikRequestJWTToken(t *testing.T) {
 				helper:       testCase.helper,
 			}
 
-			resp, err := creds.requestJWTToken()
+			resp, err := creds.requestJWTToken(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")
@@ -304,7 +305,7 @@ func TestAuthentikAuthenticate(t *testing.T) {
 			}
 			creds.jwtToken.expiresInTime = testCase.inputExpireToken
 
-			_, err := creds.Authenticate()
+			_, err := creds.Authenticate(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")

management/server/idp/azure.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -110,7 +111,7 @@ func (ac *AzureCredentials) jwtStillValid() bool {
 }
 
 // requestJWTToken performs request to get jwt token.
-func (ac *AzureCredentials) requestJWTToken() (*http.Response, error) {
+func (ac *AzureCredentials) requestJWTToken(ctx context.Context) (*http.Response, error) {
 	data := url.Values{}
 	data.Set("client_id", ac.clientConfig.ClientID)
 	data.Set("client_secret", ac.clientConfig.ClientSecret)
@@ -132,7 +133,7 @@ func (ac *AzureCredentials) requestJWTToken() (*http.Response, error) {
 	}
 	req.Header.Add("content-type", "application/x-www-form-urlencoded")
 
-	log.Debug("requesting new jwt token for azure idp manager")
+	log.WithContext(ctx).Debug("requesting new jwt token for azure idp manager")
 
 	resp, err := ac.httpClient.Do(req)
 	if err != nil {
@@ -184,7 +185,7 @@ func (ac *AzureCredentials) parseRequestJWTResponse(rawBody io.ReadCloser) (JWTT
 }
 
 // Authenticate retrieves access token to use the azure Management API.
-func (ac *AzureCredentials) Authenticate() (JWTToken, error) {
+func (ac *AzureCredentials) Authenticate(ctx context.Context) (JWTToken, error) {
 	ac.mux.Lock()
 	defer ac.mux.Unlock()
 
@@ -198,7 +199,7 @@ func (ac *AzureCredentials) Authenticate() (JWTToken, error) {
 		return ac.jwtToken, nil
 	}
 
-	resp, err := ac.requestJWTToken()
+	resp, err := ac.requestJWTToken(ctx)
 	if err != nil {
 		return ac.jwtToken, err
 	}
@@ -215,16 +216,16 @@ func (ac *AzureCredentials) Authenticate() (JWTToken, error) {
 }
 
 // CreateUser creates a new user in azure AD Idp.
-func (am *AzureManager) CreateUser(_, _, _, _ string) (*UserData, error) {
+func (am *AzureManager) CreateUser(_ context.Context, _, _, _, _ string) (*UserData, error) {
 	return nil, fmt.Errorf("method CreateUser not implemented")
 }
 
 // GetUserDataByID requests user data from keycloak via ID.
-func (am *AzureManager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
+func (am *AzureManager) GetUserDataByID(ctx context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
 	q := url.Values{}
 	q.Add("$select", profileFields)
 
-	body, err := am.get("users/"+userID, q)
+	body, err := am.get(ctx, "users/"+userID, q)
 	if err != nil {
 		return nil, err
 	}
@@ -247,11 +248,11 @@ func (am *AzureManager) GetUserDataByID(userID string, appMetadata AppMetadata)
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (am *AzureManager) GetUserByEmail(email string) ([]*UserData, error) {
+func (am *AzureManager) GetUserByEmail(ctx context.Context, email string) ([]*UserData, error) {
 	q := url.Values{}
 	q.Add("$select", profileFields)
 
-	body, err := am.get("users/"+email, q)
+	body, err := am.get(ctx, "users/"+email, q)
 	if err != nil {
 		return nil, err
 	}
@@ -273,8 +274,8 @@ func (am *AzureManager) GetUserByEmail(email string) ([]*UserData, error) {
 }
 
 // GetAccount returns all the users for a given profile.
-func (am *AzureManager) GetAccount(accountID string) ([]*UserData, error) {
-	users, err := am.getAllUsers()
+func (am *AzureManager) GetAccount(ctx context.Context, accountID string) ([]*UserData, error) {
+	users, err := am.getAllUsers(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -293,8 +294,8 @@ func (am *AzureManager) GetAccount(accountID string) ([]*UserData, error) {
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (am *AzureManager) GetAllAccounts() (map[string][]*UserData, error) {
-	users, err := am.getAllUsers()
+func (am *AzureManager) GetAllAccounts(ctx context.Context) (map[string][]*UserData, error) {
+	users, err := am.getAllUsers(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -310,19 +311,19 @@ func (am *AzureManager) GetAllAccounts() (map[string][]*UserData, error) {
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userID.
-func (am *AzureManager) UpdateUserAppMetadata(_ string, _ AppMetadata) error {
+func (am *AzureManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (am *AzureManager) InviteUserByID(_ string) error {
+func (am *AzureManager) InviteUserByID(_ context.Context, _ string) error {
 	return fmt.Errorf("method InviteUserByID not implemented")
 }
 
 // DeleteUser from Azure.
-func (am *AzureManager) DeleteUser(userID string) error {
-	jwtToken, err := am.credentials.Authenticate()
+func (am *AzureManager) DeleteUser(ctx context.Context, userID string) error {
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return err
 	}
@@ -335,7 +336,7 @@ func (am *AzureManager) DeleteUser(userID string) error {
 	req.Header.Add("authorization", "Bearer "+jwtToken.AccessToken)
 	req.Header.Add("content-type", "application/json")
 
-	log.Debugf("delete idp user %s", userID)
+	log.WithContext(ctx).Debugf("delete idp user %s", userID)
 
 	resp, err := am.httpClient.Do(req)
 	if err != nil {
@@ -358,7 +359,7 @@ func (am *AzureManager) DeleteUser(userID string) error {
 }
 
 // getAllUsers returns all users in an Azure AD account.
-func (am *AzureManager) getAllUsers() ([]*UserData, error) {
+func (am *AzureManager) getAllUsers(ctx context.Context) ([]*UserData, error) {
 	users := make([]*UserData, 0)
 
 	q := url.Values{}
@@ -366,7 +367,7 @@ func (am *AzureManager) getAllUsers() ([]*UserData, error) {
 	q.Add("$top", "500")
 
 	for nextLink := "users"; nextLink != ""; {
-		body, err := am.get(nextLink, q)
+		body, err := am.get(ctx, nextLink, q)
 		if err != nil {
 			return nil, err
 		}
@@ -391,8 +392,8 @@ func (am *AzureManager) getAllUsers() ([]*UserData, error) {
 }
 
 // get perform Get requests.
-func (am *AzureManager) get(resource string, q url.Values) ([]byte, error) {
-	jwtToken, err := am.credentials.Authenticate()
+func (am *AzureManager) get(ctx context.Context, resource string, q url.Values) ([]byte, error) {
+	jwtToken, err := am.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}

management/server/idp/azure_test.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
@@ -101,7 +102,7 @@ func TestAzureAuthenticate(t *testing.T) {
 			}
 			creds.jwtToken.expiresInTime = testCase.inputExpireToken
 
-			_, err := creds.Authenticate()
+			_, err := creds.Authenticate(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")

management/server/idp/google_workspace.go

@@ -39,12 +39,12 @@ type GoogleWorkspaceCredentials struct {
 	appMetrics   telemetry.AppMetrics
 }
 
-func (gc *GoogleWorkspaceCredentials) Authenticate() (JWTToken, error) {
+func (gc *GoogleWorkspaceCredentials) Authenticate(_ context.Context) (JWTToken, error) {
 	return JWTToken{}, nil
 }
 
 // NewGoogleWorkspaceManager creates a new instance of the GoogleWorkspaceManager.
-func NewGoogleWorkspaceManager(config GoogleWorkspaceClientConfig, appMetrics telemetry.AppMetrics) (*GoogleWorkspaceManager, error) {
+func NewGoogleWorkspaceManager(ctx context.Context, config GoogleWorkspaceClientConfig, appMetrics telemetry.AppMetrics) (*GoogleWorkspaceManager, error) {
 	httpTransport := http.DefaultTransport.(*http.Transport).Clone()
 	httpTransport.MaxIdleConns = 5
 
@@ -66,7 +66,7 @@ func NewGoogleWorkspaceManager(config GoogleWorkspaceClientConfig, appMetrics te
 	}
 
 	// Create a new Admin SDK Directory service client
-	adminCredentials, err := getGoogleCredentials(config.ServiceAccountKey)
+	adminCredentials, err := getGoogleCredentials(ctx, config.ServiceAccountKey)
 	if err != nil {
 		return nil, err
 	}
@@ -90,12 +90,12 @@ func NewGoogleWorkspaceManager(config GoogleWorkspaceClientConfig, appMetrics te
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userID and metadata map.
-func (gm *GoogleWorkspaceManager) UpdateUserAppMetadata(_ string, _ AppMetadata) error {
+func (gm *GoogleWorkspaceManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
 // GetUserDataByID requests user data from Google Workspace via ID.
-func (gm *GoogleWorkspaceManager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
+func (gm *GoogleWorkspaceManager) GetUserDataByID(_ context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
 	user, err := gm.usersService.Get(userID).Do()
 	if err != nil {
 		return nil, err
@@ -112,7 +112,7 @@ func (gm *GoogleWorkspaceManager) GetUserDataByID(userID string, appMetadata App
 }
 
 // GetAccount returns all the users for a given profile.
-func (gm *GoogleWorkspaceManager) GetAccount(accountID string) ([]*UserData, error) {
+func (gm *GoogleWorkspaceManager) GetAccount(_ context.Context, accountID string) ([]*UserData, error) {
 	users, err := gm.getAllUsers()
 	if err != nil {
 		return nil, err
@@ -132,7 +132,7 @@ func (gm *GoogleWorkspaceManager) GetAccount(accountID string) ([]*UserData, err
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (gm *GoogleWorkspaceManager) GetAllAccounts() (map[string][]*UserData, error) {
+func (gm *GoogleWorkspaceManager) GetAllAccounts(_ context.Context) (map[string][]*UserData, error) {
 	users, err := gm.getAllUsers()
 	if err != nil {
 		return nil, err
@@ -177,13 +177,13 @@ func (gm *GoogleWorkspaceManager) getAllUsers() ([]*UserData, error) {
 }
 
 // CreateUser creates a new user in Google Workspace and sends an invitation.
-func (gm *GoogleWorkspaceManager) CreateUser(_, _, _, _ string) (*UserData, error) {
+func (gm *GoogleWorkspaceManager) CreateUser(_ context.Context, _, _, _, _ string) (*UserData, error) {
 	return nil, fmt.Errorf("method CreateUser not implemented")
 }
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (gm *GoogleWorkspaceManager) GetUserByEmail(email string) ([]*UserData, error) {
+func (gm *GoogleWorkspaceManager) GetUserByEmail(_ context.Context, email string) ([]*UserData, error) {
 	user, err := gm.usersService.Get(email).Do()
 	if err != nil {
 		return nil, err
@@ -201,12 +201,12 @@ func (gm *GoogleWorkspaceManager) GetUserByEmail(email string) ([]*UserData, err
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (gm *GoogleWorkspaceManager) InviteUserByID(_ string) error {
+func (gm *GoogleWorkspaceManager) InviteUserByID(_ context.Context, _ string) error {
 	return fmt.Errorf("method InviteUserByID not implemented")
 }
 
 // DeleteUser from GoogleWorkspace.
-func (gm *GoogleWorkspaceManager) DeleteUser(userID string) error {
+func (gm *GoogleWorkspaceManager) DeleteUser(_ context.Context, userID string) error {
 	if err := gm.usersService.Delete(userID).Do(); err != nil {
 		return err
 	}
@@ -222,8 +222,8 @@ func (gm *GoogleWorkspaceManager) DeleteUser(userID string) error {
 // It decodes the base64-encoded serviceAccountKey and attempts to obtain credentials using it.
 // If that fails, it falls back to using the default Google credentials path.
 // It returns the retrieved credentials or an error if unsuccessful.
-func getGoogleCredentials(serviceAccountKey string) (*google.Credentials, error) {
-	log.Debug("retrieving google credentials from the base64 encoded service account key")
+func getGoogleCredentials(ctx context.Context, serviceAccountKey string) (*google.Credentials, error) {
+	log.WithContext(ctx).Debug("retrieving google credentials from the base64 encoded service account key")
 	decodeKey, err := base64.StdEncoding.DecodeString(serviceAccountKey)
 	if err != nil {
 		return nil, fmt.Errorf("failed to decode service account key: %w", err)
@@ -239,8 +239,8 @@ func getGoogleCredentials(serviceAccountKey string) (*google.Credentials, error)
 		return creds, nil
 	}
 
-	log.Debugf("failed to retrieve Google credentials from ServiceAccountKey: %v", err)
-	log.Debug("falling back to default google credentials location")
+	log.WithContext(ctx).Debugf("failed to retrieve Google credentials from ServiceAccountKey: %v", err)
+	log.WithContext(ctx).Debug("falling back to default google credentials location")
 
 	creds, err = google.FindDefaultCredentials(
 		context.Background(),

management/server/idp/idp.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strings"
@@ -16,14 +17,14 @@ const (
 
 // Manager idp manager interface
 type Manager interface {
-	UpdateUserAppMetadata(userId string, appMetadata AppMetadata) error
-	GetUserDataByID(userId string, appMetadata AppMetadata) (*UserData, error)
-	GetAccount(accountId string) ([]*UserData, error)
-	GetAllAccounts() (map[string][]*UserData, error)
-	CreateUser(email, name, accountID, invitedByEmail string) (*UserData, error)
-	GetUserByEmail(email string) ([]*UserData, error)
-	InviteUserByID(userID string) error
-	DeleteUser(userID string) error
+	UpdateUserAppMetadata(ctx context.Context, userId string, appMetadata AppMetadata) error
+	GetUserDataByID(ctx context.Context, userId string, appMetadata AppMetadata) (*UserData, error)
+	GetAccount(ctx context.Context, accountId string) ([]*UserData, error)
+	GetAllAccounts(ctx context.Context) (map[string][]*UserData, error)
+	CreateUser(ctx context.Context, email, name, accountID, invitedByEmail string) (*UserData, error)
+	GetUserByEmail(ctx context.Context, email string) ([]*UserData, error)
+	InviteUserByID(ctx context.Context, userID string) error
+	DeleteUser(ctx context.Context, userID string) error
 }
 
 // ClientConfig defines common client configuration for all IdP manager
@@ -51,7 +52,7 @@ type Config struct {
 
 // ManagerCredentials interface that authenticates using the credential of each type of idp
 type ManagerCredentials interface {
-	Authenticate() (JWTToken, error)
+	Authenticate(ctx context.Context) (JWTToken, error)
 }
 
 // ManagerHTTPClient http client interface for API calls
@@ -91,7 +92,7 @@ type JWTToken struct {
 }
 
 // NewManager returns a new idp manager based on the configuration that it receives
-func NewManager(config Config, appMetrics telemetry.AppMetrics) (Manager, error) {
+func NewManager(ctx context.Context, config Config, appMetrics telemetry.AppMetrics) (Manager, error) {
 	if config.ClientConfig != nil {
 		config.ClientConfig.Issuer = strings.TrimSuffix(config.ClientConfig.Issuer, "/")
 	}
@@ -175,7 +176,7 @@ func NewManager(config Config, appMetrics telemetry.AppMetrics) (Manager, error)
 			ServiceAccountKey: config.ExtraConfig["ServiceAccountKey"],
 			CustomerID:        config.ExtraConfig["CustomerId"],
 		}
-		return NewGoogleWorkspaceManager(googleClientConfig, appMetrics)
+		return NewGoogleWorkspaceManager(ctx, googleClientConfig, appMetrics)
 	case "jumpcloud":
 		jumpcloudConfig := JumpCloudClientConfig{
 			APIToken: config.ExtraConfig["ApiToken"],

management/server/idp/jumpcloud.go

@@ -74,7 +74,7 @@ func NewJumpCloudManager(config JumpCloudClientConfig, appMetrics telemetry.AppM
 }
 
 // Authenticate retrieves access token to use the JumpCloud user API.
-func (jc *JumpCloudCredentials) Authenticate() (JWTToken, error) {
+func (jc *JumpCloudCredentials) Authenticate(_ context.Context) (JWTToken, error) {
 	return JWTToken{}, nil
 }
 
@@ -85,12 +85,12 @@ func (jm *JumpCloudManager) authenticationContext() context.Context {
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userID and metadata map.
-func (jm *JumpCloudManager) UpdateUserAppMetadata(_ string, _ AppMetadata) error {
+func (jm *JumpCloudManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
 // GetUserDataByID requests user data from JumpCloud via ID.
-func (jm *JumpCloudManager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
+func (jm *JumpCloudManager) GetUserDataByID(_ context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
 	authCtx := jm.authenticationContext()
 	user, resp, err := jm.client.SystemusersApi.SystemusersGet(authCtx, userID, contentType, accept, nil)
 	if err != nil {
@@ -116,7 +116,7 @@ func (jm *JumpCloudManager) GetUserDataByID(userID string, appMetadata AppMetada
 }
 
 // GetAccount returns all the users for a given profile.
-func (jm *JumpCloudManager) GetAccount(accountID string) ([]*UserData, error) {
+func (jm *JumpCloudManager) GetAccount(_ context.Context, accountID string) ([]*UserData, error) {
 	authCtx := jm.authenticationContext()
 	userList, resp, err := jm.client.SearchApi.SearchSystemusersPost(authCtx, contentType, accept, nil)
 	if err != nil {
@@ -148,7 +148,7 @@ func (jm *JumpCloudManager) GetAccount(accountID string) ([]*UserData, error) {
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (jm *JumpCloudManager) GetAllAccounts() (map[string][]*UserData, error) {
+func (jm *JumpCloudManager) GetAllAccounts(_ context.Context) (map[string][]*UserData, error) {
 	authCtx := jm.authenticationContext()
 	userList, resp, err := jm.client.SearchApi.SearchSystemusersPost(authCtx, contentType, accept, nil)
 	if err != nil {
@@ -177,13 +177,13 @@ func (jm *JumpCloudManager) GetAllAccounts() (map[string][]*UserData, error) {
 }
 
 // CreateUser creates a new user in JumpCloud Idp and sends an invitation.
-func (jm *JumpCloudManager) CreateUser(_, _, _, _ string) (*UserData, error) {
+func (jm *JumpCloudManager) CreateUser(_ context.Context, _, _, _, _ string) (*UserData, error) {
 	return nil, fmt.Errorf("method CreateUser not implemented")
 }
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (jm *JumpCloudManager) GetUserByEmail(email string) ([]*UserData, error) {
+func (jm *JumpCloudManager) GetUserByEmail(_ context.Context, email string) ([]*UserData, error) {
 	searchFilter := map[string]interface{}{
 		"searchFilter": map[string]interface{}{
 			"filter": []string{email},
@@ -219,12 +219,12 @@ func (jm *JumpCloudManager) GetUserByEmail(email string) ([]*UserData, error) {
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (jm *JumpCloudManager) InviteUserByID(_ string) error {
+func (jm *JumpCloudManager) InviteUserByID(_ context.Context, _ string) error {
 	return fmt.Errorf("method InviteUserByID not implemented")
 }
 
 // DeleteUser from jumpCloud directory
-func (jm *JumpCloudManager) DeleteUser(userID string) error {
+func (jm *JumpCloudManager) DeleteUser(_ context.Context, userID string) error {
 	authCtx := jm.authenticationContext()
 	_, resp, err := jm.client.SystemusersApi.SystemusersDelete(authCtx, userID, contentType, accept, nil)
 	if err != nil {

management/server/idp/keycloak.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -109,7 +110,7 @@ func (kc *KeycloakCredentials) jwtStillValid() bool {
 }
 
 // requestJWTToken performs request to get jwt token.
-func (kc *KeycloakCredentials) requestJWTToken() (*http.Response, error) {
+func (kc *KeycloakCredentials) requestJWTToken(ctx context.Context) (*http.Response, error) {
 	data := url.Values{}
 	data.Set("client_id", kc.clientConfig.ClientID)
 	data.Set("client_secret", kc.clientConfig.ClientSecret)
@@ -122,7 +123,7 @@ func (kc *KeycloakCredentials) requestJWTToken() (*http.Response, error) {
 	}
 	req.Header.Add("content-type", "application/x-www-form-urlencoded")
 
-	log.Debug("requesting new jwt token for keycloak idp manager")
+	log.WithContext(ctx).Debug("requesting new jwt token for keycloak idp manager")
 
 	resp, err := kc.httpClient.Do(req)
 	if err != nil {
@@ -174,7 +175,7 @@ func (kc *KeycloakCredentials) parseRequestJWTResponse(rawBody io.ReadCloser) (J
 }
 
 // Authenticate retrieves access token to use the keycloak Management API.
-func (kc *KeycloakCredentials) Authenticate() (JWTToken, error) {
+func (kc *KeycloakCredentials) Authenticate(ctx context.Context) (JWTToken, error) {
 	kc.mux.Lock()
 	defer kc.mux.Unlock()
 
@@ -188,7 +189,7 @@ func (kc *KeycloakCredentials) Authenticate() (JWTToken, error) {
 		return kc.jwtToken, nil
 	}
 
-	resp, err := kc.requestJWTToken()
+	resp, err := kc.requestJWTToken(ctx)
 	if err != nil {
 		return kc.jwtToken, err
 	}
@@ -205,18 +206,18 @@ func (kc *KeycloakCredentials) Authenticate() (JWTToken, error) {
 }
 
 // CreateUser creates a new user in keycloak Idp and sends an invite.
-func (km *KeycloakManager) CreateUser(_, _, _, _ string) (*UserData, error) {
+func (km *KeycloakManager) CreateUser(_ context.Context, _, _, _, _ string) (*UserData, error) {
 	return nil, fmt.Errorf("method CreateUser not implemented")
 }
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (km *KeycloakManager) GetUserByEmail(email string) ([]*UserData, error) {
+func (km *KeycloakManager) GetUserByEmail(ctx context.Context, email string) ([]*UserData, error) {
 	q := url.Values{}
 	q.Add("email", email)
 	q.Add("exact", "true")
 
-	body, err := km.get("users", q)
+	body, err := km.get(ctx, "users", q)
 	if err != nil {
 		return nil, err
 	}
@@ -240,8 +241,8 @@ func (km *KeycloakManager) GetUserByEmail(email string) ([]*UserData, error) {
 }
 
 // GetUserDataByID requests user data from keycloak via ID.
-func (km *KeycloakManager) GetUserDataByID(userID string, _ AppMetadata) (*UserData, error) {
-	body, err := km.get("users/"+userID, nil)
+func (km *KeycloakManager) GetUserDataByID(ctx context.Context, userID string, _ AppMetadata) (*UserData, error) {
+	body, err := km.get(ctx, "users/"+userID, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -260,8 +261,8 @@ func (km *KeycloakManager) GetUserDataByID(userID string, _ AppMetadata) (*UserD
 }
 
 // GetAccount returns all the users for a given account profile.
-func (km *KeycloakManager) GetAccount(accountID string) ([]*UserData, error) {
-	profiles, err := km.fetchAllUserProfiles()
+func (km *KeycloakManager) GetAccount(ctx context.Context, accountID string) ([]*UserData, error) {
+	profiles, err := km.fetchAllUserProfiles(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -283,8 +284,8 @@ func (km *KeycloakManager) GetAccount(accountID string) ([]*UserData, error) {
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (km *KeycloakManager) GetAllAccounts() (map[string][]*UserData, error) {
-	profiles, err := km.fetchAllUserProfiles()
+func (km *KeycloakManager) GetAllAccounts(ctx context.Context) (map[string][]*UserData, error) {
+	profiles, err := km.fetchAllUserProfiles(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -303,19 +304,19 @@ func (km *KeycloakManager) GetAllAccounts() (map[string][]*UserData, error) {
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userID and metadata map.
-func (km *KeycloakManager) UpdateUserAppMetadata(_ string, _ AppMetadata) error {
+func (km *KeycloakManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (km *KeycloakManager) InviteUserByID(_ string) error {
+func (km *KeycloakManager) InviteUserByID(_ context.Context, _ string) error {
 	return fmt.Errorf("method InviteUserByID not implemented")
 }
 
 // DeleteUser from Keycloak by user ID.
-func (km *KeycloakManager) DeleteUser(userID string) error {
-	jwtToken, err := km.credentials.Authenticate()
+func (km *KeycloakManager) DeleteUser(ctx context.Context, userID string) error {
+	jwtToken, err := km.credentials.Authenticate(ctx)
 	if err != nil {
 		return err
 	}
@@ -353,8 +354,8 @@ func (km *KeycloakManager) DeleteUser(userID string) error {
 	return nil
 }
 
-func (km *KeycloakManager) fetchAllUserProfiles() ([]keycloakProfile, error) {
-	totalUsers, err := km.totalUsersCount()
+func (km *KeycloakManager) fetchAllUserProfiles(ctx context.Context) ([]keycloakProfile, error) {
+	totalUsers, err := km.totalUsersCount(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -362,7 +363,7 @@ func (km *KeycloakManager) fetchAllUserProfiles() ([]keycloakProfile, error) {
 	q := url.Values{}
 	q.Add("max", fmt.Sprint(*totalUsers))
 
-	body, err := km.get("users", q)
+	body, err := km.get(ctx, "users", q)
 	if err != nil {
 		return nil, err
 	}
@@ -377,8 +378,8 @@ func (km *KeycloakManager) fetchAllUserProfiles() ([]keycloakProfile, error) {
 }
 
 // get perform Get requests.
-func (km *KeycloakManager) get(resource string, q url.Values) ([]byte, error) {
-	jwtToken, err := km.credentials.Authenticate()
+func (km *KeycloakManager) get(ctx context.Context, resource string, q url.Values) ([]byte, error) {
+	jwtToken, err := km.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -414,8 +415,8 @@ func (km *KeycloakManager) get(resource string, q url.Values) ([]byte, error) {
 
 // totalUsersCount returns the total count of all user created.
 // Used when fetching all registered accounts with pagination.
-func (km *KeycloakManager) totalUsersCount() (*int, error) {
-	body, err := km.get("users/count", nil)
+func (km *KeycloakManager) totalUsersCount(ctx context.Context) (*int, error) {
+	body, err := km.get(ctx, "users/count", nil)
 	if err != nil {
 		return nil, err
 	}

management/server/idp/keycloak_test.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"strings"
@@ -128,7 +129,7 @@ func TestKeycloakRequestJWTToken(t *testing.T) {
 				helper:       testCase.helper,
 			}
 
-			resp, err := creds.requestJWTToken()
+			resp, err := creds.requestJWTToken(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")
@@ -294,7 +295,7 @@ func TestKeycloakAuthenticate(t *testing.T) {
 			}
 			creds.jwtToken.expiresInTime = testCase.inputExpireToken
 
-			_, err := creds.Authenticate()
+			_, err := creds.Authenticate(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")

management/server/idp/mock.go

@@ -1,77 +1,79 @@
 package idp
 
+import "context"
+
 // MockIDP is a mock implementation of the IDP interface
 type MockIDP struct {
-	UpdateUserAppMetadataFunc func(userId string, appMetadata AppMetadata) error
-	GetUserDataByIDFunc       func(userId string, appMetadata AppMetadata) (*UserData, error)
-	GetAccountFunc            func(accountId string) ([]*UserData, error)
-	GetAllAccountsFunc        func() (map[string][]*UserData, error)
-	CreateUserFunc            func(email, name, accountID, invitedByEmail string) (*UserData, error)
-	GetUserByEmailFunc        func(email string) ([]*UserData, error)
-	InviteUserByIDFunc        func(userID string) error
-	DeleteUserFunc            func(userID string) error
+	UpdateUserAppMetadataFunc func(ctx context.Context, userId string, appMetadata AppMetadata) error
+	GetUserDataByIDFunc       func(ctx context.Context, userId string, appMetadata AppMetadata) (*UserData, error)
+	GetAccountFunc            func(ctx context.Context, accountId string) ([]*UserData, error)
+	GetAllAccountsFunc        func(ctx context.Context) (map[string][]*UserData, error)
+	CreateUserFunc            func(ctx context.Context, email, name, accountID, invitedByEmail string) (*UserData, error)
+	GetUserByEmailFunc        func(ctx context.Context, email string) ([]*UserData, error)
+	InviteUserByIDFunc        func(ctx context.Context, userID string) error
+	DeleteUserFunc            func(ctx context.Context, userID string) error
 }
 
 // UpdateUserAppMetadata is a mock implementation of the IDP interface UpdateUserAppMetadata method
-func (m *MockIDP) UpdateUserAppMetadata(userId string, appMetadata AppMetadata) error {
+func (m *MockIDP) UpdateUserAppMetadata(ctx context.Context, userId string, appMetadata AppMetadata) error {
 	if m.UpdateUserAppMetadataFunc != nil {
-		return m.UpdateUserAppMetadataFunc(userId, appMetadata)
+		return m.UpdateUserAppMetadataFunc(ctx, userId, appMetadata)
 	}
 	return nil
 }
 
 // GetUserDataByID is a mock implementation of the IDP interface GetUserDataByID method
-func (m *MockIDP) GetUserDataByID(userId string, appMetadata AppMetadata) (*UserData, error) {
+func (m *MockIDP) GetUserDataByID(ctx context.Context, userId string, appMetadata AppMetadata) (*UserData, error) {
 	if m.GetUserDataByIDFunc != nil {
-		return m.GetUserDataByIDFunc(userId, appMetadata)
+		return m.GetUserDataByIDFunc(ctx, userId, appMetadata)
 	}
 	return nil, nil
 }
 
 // GetAccount is a mock implementation of the IDP interface GetAccount method
-func (m *MockIDP) GetAccount(accountId string) ([]*UserData, error) {
+func (m *MockIDP) GetAccount(ctx context.Context, accountId string) ([]*UserData, error) {
 	if m.GetAccountFunc != nil {
-		return m.GetAccountFunc(accountId)
+		return m.GetAccountFunc(ctx, accountId)
 	}
 	return nil, nil
 }
 
 // GetAllAccounts is a mock implementation of the IDP interface GetAllAccounts method
-func (m *MockIDP) GetAllAccounts() (map[string][]*UserData, error) {
+func (m *MockIDP) GetAllAccounts(ctx context.Context) (map[string][]*UserData, error) {
 	if m.GetAllAccountsFunc != nil {
-		return m.GetAllAccountsFunc()
+		return m.GetAllAccountsFunc(ctx)
 	}
 	return nil, nil
 }
 
 // CreateUser is a mock implementation of the IDP interface CreateUser method
-func (m *MockIDP) CreateUser(email, name, accountID, invitedByEmail string) (*UserData, error) {
+func (m *MockIDP) CreateUser(ctx context.Context, email, name, accountID, invitedByEmail string) (*UserData, error) {
 	if m.CreateUserFunc != nil {
-		return m.CreateUserFunc(email, name, accountID, invitedByEmail)
+		return m.CreateUserFunc(ctx, email, name, accountID, invitedByEmail)
 	}
 	return nil, nil
 }
 
 // GetUserByEmail is a mock implementation of the IDP interface GetUserByEmail method
-func (m *MockIDP) GetUserByEmail(email string) ([]*UserData, error) {
+func (m *MockIDP) GetUserByEmail(ctx context.Context, email string) ([]*UserData, error) {
 	if m.GetUserByEmailFunc != nil {
-		return m.GetUserByEmailFunc(email)
+		return m.GetUserByEmailFunc(ctx, email)
 	}
 	return nil, nil
 }
 
 // InviteUserByID is a mock implementation of the IDP interface InviteUserByID method
-func (m *MockIDP) InviteUserByID(userID string) error {
+func (m *MockIDP) InviteUserByID(ctx context.Context, userID string) error {
 	if m.InviteUserByIDFunc != nil {
-		return m.InviteUserByIDFunc(userID)
+		return m.InviteUserByIDFunc(ctx, userID)
 	}
 	return nil
 }
 
 // DeleteUser is a mock implementation of the IDP interface DeleteUser method
-func (m *MockIDP) DeleteUser(userID string) error {
+func (m *MockIDP) DeleteUser(ctx context.Context, userID string) error {
 	if m.DeleteUserFunc != nil {
-		return m.DeleteUserFunc(userID)
+		return m.DeleteUserFunc(ctx, userID)
 	}
 	return nil
 }

management/server/idp/okta.go

@@ -94,17 +94,17 @@ func NewOktaManager(config OktaClientConfig, appMetrics telemetry.AppMetrics) (*
 }
 
 // Authenticate retrieves access token to use the okta user API.
-func (oc *OktaCredentials) Authenticate() (JWTToken, error) {
+func (oc *OktaCredentials) Authenticate(_ context.Context) (JWTToken, error) {
 	return JWTToken{}, nil
 }
 
 // CreateUser creates a new user in okta Idp and sends an invitation.
-func (om *OktaManager) CreateUser(_, _, _, _ string) (*UserData, error) {
+func (om *OktaManager) CreateUser(_ context.Context, _, _, _, _ string) (*UserData, error) {
 	return nil, fmt.Errorf("method CreateUser not implemented")
 }
 
 // GetUserDataByID requests user data from keycloak via ID.
-func (om *OktaManager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
+func (om *OktaManager) GetUserDataByID(_ context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
 	user, resp, err := om.client.User.GetUser(context.Background(), userID)
 	if err != nil {
 		return nil, err
@@ -132,7 +132,7 @@ func (om *OktaManager) GetUserDataByID(userID string, appMetadata AppMetadata) (
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (om *OktaManager) GetUserByEmail(email string) ([]*UserData, error) {
+func (om *OktaManager) GetUserByEmail(_ context.Context, email string) ([]*UserData, error) {
 	user, resp, err := om.client.User.GetUser(context.Background(), url.QueryEscape(email))
 	if err != nil {
 		return nil, err
@@ -160,7 +160,7 @@ func (om *OktaManager) GetUserByEmail(email string) ([]*UserData, error) {
 }
 
 // GetAccount returns all the users for a given profile.
-func (om *OktaManager) GetAccount(accountID string) ([]*UserData, error) {
+func (om *OktaManager) GetAccount(_ context.Context, accountID string) ([]*UserData, error) {
 	users, err := om.getAllUsers()
 	if err != nil {
 		return nil, err
@@ -180,7 +180,7 @@ func (om *OktaManager) GetAccount(accountID string) ([]*UserData, error) {
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (om *OktaManager) GetAllAccounts() (map[string][]*UserData, error) {
+func (om *OktaManager) GetAllAccounts(_ context.Context) (map[string][]*UserData, error) {
 	users, err := om.getAllUsers()
 	if err != nil {
 		return nil, err
@@ -242,18 +242,18 @@ func (om *OktaManager) getAllUsers() ([]*UserData, error) {
 }
 
 // UpdateUserAppMetadata updates user app metadata based on userID and metadata map.
-func (om *OktaManager) UpdateUserAppMetadata(userID string, appMetadata AppMetadata) error {
+func (om *OktaManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (om *OktaManager) InviteUserByID(_ string) error {
+func (om *OktaManager) InviteUserByID(_ context.Context, _ string) error {
 	return fmt.Errorf("method InviteUserByID not implemented")
 }
 
 // DeleteUser from Okta
-func (om *OktaManager) DeleteUser(userID string) error {
+func (om *OktaManager) DeleteUser(_ context.Context, userID string) error {
 	resp, err := om.client.User.DeactivateOrDeleteUser(context.Background(), userID, nil)
 	if err != nil {
 		return err

management/server/idp/zitadel.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -149,7 +150,7 @@ func (zc *ZitadelCredentials) jwtStillValid() bool {
 }
 
 // requestJWTToken performs request to get jwt token.
-func (zc *ZitadelCredentials) requestJWTToken() (*http.Response, error) {
+func (zc *ZitadelCredentials) requestJWTToken(ctx context.Context) (*http.Response, error) {
 	data := url.Values{}
 	data.Set("client_id", zc.clientConfig.ClientID)
 	data.Set("client_secret", zc.clientConfig.ClientSecret)
@@ -163,7 +164,7 @@ func (zc *ZitadelCredentials) requestJWTToken() (*http.Response, error) {
 	}
 	req.Header.Add("content-type", "application/x-www-form-urlencoded")
 
-	log.Debug("requesting new jwt token for zitadel idp manager")
+	log.WithContext(ctx).Debug("requesting new jwt token for zitadel idp manager")
 
 	resp, err := zc.httpClient.Do(req)
 	if err != nil {
@@ -215,7 +216,7 @@ func (zc *ZitadelCredentials) parseRequestJWTResponse(rawBody io.ReadCloser) (JW
 }
 
 // Authenticate retrieves access token to use the Zitadel Management API.
-func (zc *ZitadelCredentials) Authenticate() (JWTToken, error) {
+func (zc *ZitadelCredentials) Authenticate(ctx context.Context) (JWTToken, error) {
 	zc.mux.Lock()
 	defer zc.mux.Unlock()
 
@@ -229,7 +230,7 @@ func (zc *ZitadelCredentials) Authenticate() (JWTToken, error) {
 		return zc.jwtToken, nil
 	}
 
-	resp, err := zc.requestJWTToken()
+	resp, err := zc.requestJWTToken(ctx)
 	if err != nil {
 		return zc.jwtToken, err
 	}
@@ -246,7 +247,7 @@ func (zc *ZitadelCredentials) Authenticate() (JWTToken, error) {
 }
 
 // CreateUser creates a new user in zitadel Idp and sends an invite via Zitadel.
-func (zm *ZitadelManager) CreateUser(email, name, accountID, invitedByEmail string) (*UserData, error) {
+func (zm *ZitadelManager) CreateUser(ctx context.Context, email, name, accountID, invitedByEmail string) (*UserData, error) {
 	firstLast := strings.SplitN(name, " ", 2)
 
 	var addUser = map[string]any{
@@ -269,7 +270,7 @@ func (zm *ZitadelManager) CreateUser(email, name, accountID, invitedByEmail stri
 		return nil, err
 	}
 
-	body, err := zm.post("users/human/_import", string(payload))
+	body, err := zm.post(ctx, "users/human/_import", string(payload))
 	if err != nil {
 		return nil, err
 	}
@@ -300,7 +301,7 @@ func (zm *ZitadelManager) CreateUser(email, name, accountID, invitedByEmail stri
 
 // GetUserByEmail searches users with a given email.
 // If no users have been found, this function returns an empty list.
-func (zm *ZitadelManager) GetUserByEmail(email string) ([]*UserData, error) {
+func (zm *ZitadelManager) GetUserByEmail(ctx context.Context, email string) ([]*UserData, error) {
 	searchByEmail := zitadelAttributes{
 		"queries": {
 			{
@@ -316,7 +317,7 @@ func (zm *ZitadelManager) GetUserByEmail(email string) ([]*UserData, error) {
 		return nil, err
 	}
 
-	body, err := zm.post("users/_search", string(payload))
+	body, err := zm.post(ctx, "users/_search", string(payload))
 	if err != nil {
 		return nil, err
 	}
@@ -340,8 +341,8 @@ func (zm *ZitadelManager) GetUserByEmail(email string) ([]*UserData, error) {
 }
 
 // GetUserDataByID requests user data from zitadel via ID.
-func (zm *ZitadelManager) GetUserDataByID(userID string, appMetadata AppMetadata) (*UserData, error) {
-	body, err := zm.get("users/"+userID, nil)
+func (zm *ZitadelManager) GetUserDataByID(ctx context.Context, userID string, appMetadata AppMetadata) (*UserData, error) {
+	body, err := zm.get(ctx, "users/"+userID, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -363,8 +364,8 @@ func (zm *ZitadelManager) GetUserDataByID(userID string, appMetadata AppMetadata
 }
 
 // GetAccount returns all the users for a given profile.
-func (zm *ZitadelManager) GetAccount(accountID string) ([]*UserData, error) {
-	body, err := zm.post("users/_search", "")
+func (zm *ZitadelManager) GetAccount(ctx context.Context, accountID string) ([]*UserData, error) {
+	body, err := zm.post(ctx, "users/_search", "")
 	if err != nil {
 		return nil, err
 	}
@@ -392,8 +393,8 @@ func (zm *ZitadelManager) GetAccount(accountID string) ([]*UserData, error) {
 
 // GetAllAccounts gets all registered accounts with corresponding user data.
 // It returns a list of users indexed by accountID.
-func (zm *ZitadelManager) GetAllAccounts() (map[string][]*UserData, error) {
-	body, err := zm.post("users/_search", "")
+func (zm *ZitadelManager) GetAllAccounts(ctx context.Context) (map[string][]*UserData, error) {
+	body, err := zm.post(ctx, "users/_search", "")
 	if err != nil {
 		return nil, err
 	}
@@ -419,7 +420,7 @@ func (zm *ZitadelManager) GetAllAccounts() (map[string][]*UserData, error) {
 
 // UpdateUserAppMetadata updates user app metadata based on userID and metadata map.
 // Metadata values are base64 encoded.
-func (zm *ZitadelManager) UpdateUserAppMetadata(_ string, _ AppMetadata) error {
+func (zm *ZitadelManager) UpdateUserAppMetadata(_ context.Context, _ string, _ AppMetadata) error {
 	return nil
 }
 
@@ -429,7 +430,7 @@ type inviteUserRequest struct {
 
 // InviteUserByID resend invitations to users who haven't activated,
 // their accounts prior to the expiration period.
-func (zm *ZitadelManager) InviteUserByID(userID string) error {
+func (zm *ZitadelManager) InviteUserByID(ctx context.Context, userID string) error {
 	inviteUser := inviteUserRequest{
 		Email: userID,
 	}
@@ -440,14 +441,14 @@ func (zm *ZitadelManager) InviteUserByID(userID string) error {
 	}
 
 	// don't care about the body in the response
-	_, err = zm.post(fmt.Sprintf("users/%s/_resend_initialization", userID), string(payload))
+	_, err = zm.post(ctx, fmt.Sprintf("users/%s/_resend_initialization", userID), string(payload))
 	return err
 }
 
 // DeleteUser from Zitadel
-func (zm *ZitadelManager) DeleteUser(userID string) error {
+func (zm *ZitadelManager) DeleteUser(ctx context.Context, userID string) error {
 	resource := fmt.Sprintf("users/%s", userID)
-	if err := zm.delete(resource); err != nil {
+	if err := zm.delete(ctx, resource); err != nil {
 		return err
 	}
 
@@ -459,8 +460,8 @@ func (zm *ZitadelManager) DeleteUser(userID string) error {
 }
 
 // post perform Post requests.
-func (zm *ZitadelManager) post(resource string, body string) ([]byte, error) {
-	jwtToken, err := zm.credentials.Authenticate()
+func (zm *ZitadelManager) post(ctx context.Context, resource string, body string) ([]byte, error) {
+	jwtToken, err := zm.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -495,8 +496,8 @@ func (zm *ZitadelManager) post(resource string, body string) ([]byte, error) {
 }
 
 // delete perform Delete requests.
-func (zm *ZitadelManager) delete(resource string) error {
-	jwtToken, err := zm.credentials.Authenticate()
+func (zm *ZitadelManager) delete(ctx context.Context, resource string) error {
+	jwtToken, err := zm.credentials.Authenticate(ctx)
 	if err != nil {
 		return err
 	}
@@ -531,8 +532,8 @@ func (zm *ZitadelManager) delete(resource string) error {
 }
 
 // get perform Get requests.
-func (zm *ZitadelManager) get(resource string, q url.Values) ([]byte, error) {
-	jwtToken, err := zm.credentials.Authenticate()
+func (zm *ZitadelManager) get(ctx context.Context, resource string, q url.Values) ([]byte, error) {
+	jwtToken, err := zm.credentials.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}

management/server/idp/zitadel_test.go

@@ -1,6 +1,7 @@
 package idp
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"strings"
@@ -108,7 +109,7 @@ func TestZitadelRequestJWTToken(t *testing.T) {
 				helper:       testCase.helper,
 			}
 
-			resp, err := creds.requestJWTToken()
+			resp, err := creds.requestJWTToken(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")
@@ -274,7 +275,7 @@ func TestZitadelAuthenticate(t *testing.T) {
 			}
 			creds.jwtToken.expiresInTime = testCase.inputExpireToken
 
-			_, err := creds.Authenticate()
+			_, err := creds.Authenticate(context.Background())
 			if err != nil {
 				if testCase.expectedFuncExitErrDiff != nil {
 					assert.EqualError(t, err, testCase.expectedFuncExitErrDiff.Error(), "errors should be the same")