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/http/accounts_handler.go

@@ -35,34 +35,34 @@ func NewAccountsHandler(accountManager server.AccountManager, authCfg AuthCfg) *
 // GetAllAccounts is HTTP GET handler that returns a list of accounts. Effectively returns just a single account.
 func (h *AccountsHandler) GetAllAccounts(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	if !(user.HasAdminPower() || user.IsServiceUser) {
-		util.WriteError(status.Errorf(status.PermissionDenied, "the user has no permission to access account data"), w)
+		util.WriteError(r.Context(), status.Errorf(status.PermissionDenied, "the user has no permission to access account data"), w)
 		return
 	}
 
 	resp := toAccountResponse(account)
-	util.WriteJSONObject(w, []*api.Account{resp})
+	util.WriteJSONObject(r.Context(), w, []*api.Account{resp})
 }
 
 // UpdateAccount is HTTP PUT handler that updates the provided account. Updates only account settings (server.Settings)
 func (h *AccountsHandler) UpdateAccount(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	_, user, err := h.accountManager.GetAccountFromToken(claims)
+	_, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	accountID := vars["accountId"]
 	if len(accountID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid accountID ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid accountID ID"), w)
 		return
 	}
 
@@ -96,15 +96,15 @@ func (h *AccountsHandler) UpdateAccount(w http.ResponseWriter, r *http.Request)
 		settings.JWTAllowGroups = *req.Settings.JwtAllowGroups
 	}
 
-	updatedAccount, err := h.accountManager.UpdateAccountSettings(accountID, user.Id, settings)
+	updatedAccount, err := h.accountManager.UpdateAccountSettings(r.Context(), accountID, user.Id, settings)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	resp := toAccountResponse(updatedAccount)
 
-	util.WriteJSONObject(w, &resp)
+	util.WriteJSONObject(r.Context(), w, &resp)
 }
 
 // DeleteAccount is a HTTP DELETE handler to delete an account
@@ -118,17 +118,17 @@ func (h *AccountsHandler) DeleteAccount(w http.ResponseWriter, r *http.Request)
 	vars := mux.Vars(r)
 	targetAccountID := vars["accountId"]
 	if len(targetAccountID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid account ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid account ID"), w)
 		return
 	}
 
-	err := h.accountManager.DeleteAccount(targetAccountID, claims.UserId)
+	err := h.accountManager.DeleteAccount(r.Context(), targetAccountID, claims.UserId)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 func toAccountResponse(account *server.Account) *api.Account {

management/server/http/accounts_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -22,10 +23,10 @@ import (
 func initAccountsTestData(account *server.Account, admin *server.User) *AccountsHandler {
 	return &AccountsHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(ctx context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return account, admin, nil
 			},
-			UpdateAccountSettingsFunc: func(accountID, userID string, newSettings *server.Settings) (*server.Account, error) {
+			UpdateAccountSettingsFunc: func(ctx context.Context, accountID, userID string, newSettings *server.Settings) (*server.Account, error) {
 				halfYearLimit := 180 * 24 * time.Hour
 				if newSettings.PeerLoginExpiration > halfYearLimit {
 					return nil, status.Errorf(status.InvalidArgument, "peer login expiration can't be larger than 180 days")

management/server/http/dns_settings_handler.go

@@ -32,16 +32,16 @@ func NewDNSSettingsHandler(accountManager server.AccountManager, authCfg AuthCfg
 // GetDNSSettings returns the DNS settings for the account
 func (h *DNSSettingsHandler) GetDNSSettings(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		log.Error(err)
+		log.WithContext(r.Context()).Error(err)
 		http.Redirect(w, r, "/", http.StatusInternalServerError)
 		return
 	}
 
-	dnsSettings, err := h.accountManager.GetDNSSettings(account.Id, user.Id)
+	dnsSettings, err := h.accountManager.GetDNSSettings(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -49,15 +49,15 @@ func (h *DNSSettingsHandler) GetDNSSettings(w http.ResponseWriter, r *http.Reque
 		DisabledManagementGroups: dnsSettings.DisabledManagementGroups,
 	}
 
-	util.WriteJSONObject(w, apiDNSSettings)
+	util.WriteJSONObject(r.Context(), w, apiDNSSettings)
 }
 
 // UpdateDNSSettings handles update to DNS settings of an account
 func (h *DNSSettingsHandler) UpdateDNSSettings(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -72,9 +72,9 @@ func (h *DNSSettingsHandler) UpdateDNSSettings(w http.ResponseWriter, r *http.Re
 		DisabledManagementGroups: req.DisabledManagementGroups,
 	}
 
-	err = h.accountManager.SaveDNSSettings(account.Id, user.Id, updateDNSSettings)
+	err = h.accountManager.SaveDNSSettings(r.Context(), account.Id, user.Id, updateDNSSettings)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -82,5 +82,5 @@ func (h *DNSSettingsHandler) UpdateDNSSettings(w http.ResponseWriter, r *http.Re
 		DisabledManagementGroups: updateDNSSettings.DisabledManagementGroups,
 	}
 
-	util.WriteJSONObject(w, &resp)
+	util.WriteJSONObject(r.Context(), w, &resp)
 }

management/server/http/dns_settings_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -42,16 +43,16 @@ var testingDNSSettingsAccount = &server.Account{
 func initDNSSettingsTestData() *DNSSettingsHandler {
 	return &DNSSettingsHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetDNSSettingsFunc: func(accountID string, userID string) (*server.DNSSettings, error) {
+			GetDNSSettingsFunc: func(ctx context.Context, accountID string, userID string) (*server.DNSSettings, error) {
 				return &testingDNSSettingsAccount.DNSSettings, nil
 			},
-			SaveDNSSettingsFunc: func(accountID string, userID string, dnsSettingsToSave *server.DNSSettings) error {
+			SaveDNSSettingsFunc: func(ctx context.Context, accountID string, userID string, dnsSettingsToSave *server.DNSSettings) error {
 				if dnsSettingsToSave != nil {
 					return nil
 				}
 				return status.Errorf(status.InvalidArgument, "the dns settings provided are nil")
 			},
-			GetAccountFromTokenFunc: func(_ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(ctx context.Context, _ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return testingDNSSettingsAccount, testingDNSSettingsAccount.Users[testDNSSettingsUserID], nil
 			},
 		},

management/server/http/events_handler.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
@@ -33,16 +34,16 @@ func NewEventsHandler(accountManager server.AccountManager, authCfg AuthCfg) *Ev
 // GetAllEvents list of the given account
 func (h *EventsHandler) GetAllEvents(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		log.Error(err)
+		log.WithContext(r.Context()).Error(err)
 		http.Redirect(w, r, "/", http.StatusInternalServerError)
 		return
 	}
 
-	accountEvents, err := h.accountManager.GetEvents(account.Id, user.Id)
+	accountEvents, err := h.accountManager.GetEvents(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 	events := make([]*api.Event, len(accountEvents))
@@ -50,20 +51,20 @@ func (h *EventsHandler) GetAllEvents(w http.ResponseWriter, r *http.Request) {
 		events[i] = toEventResponse(e)
 	}
 
-	err = h.fillEventsWithUserInfo(events, account.Id, user.Id)
+	err = h.fillEventsWithUserInfo(r.Context(), events, account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, events)
+	util.WriteJSONObject(r.Context(), w, events)
 }
 
-func (h *EventsHandler) fillEventsWithUserInfo(events []*api.Event, accountId, userId string) error {
+func (h *EventsHandler) fillEventsWithUserInfo(ctx context.Context, events []*api.Event, accountId, userId string) error {
 	// build email, name maps based on users
-	userInfos, err := h.accountManager.GetUsersFromAccount(accountId, userId)
+	userInfos, err := h.accountManager.GetUsersFromAccount(ctx, accountId, userId)
 	if err != nil {
-		log.Errorf("failed to get users from account: %s", err)
+		log.WithContext(ctx).Errorf("failed to get users from account: %s", err)
 		return err
 	}
 
@@ -80,7 +81,7 @@ func (h *EventsHandler) fillEventsWithUserInfo(events []*api.Event, accountId, u
 		if event.InitiatorEmail == "" {
 			event.InitiatorEmail, ok = emails[event.InitiatorId]
 			if !ok {
-				log.Warnf("failed to resolve email for initiator: %s", event.InitiatorId)
+				log.WithContext(ctx).Warnf("failed to resolve email for initiator: %s", event.InitiatorId)
 			}
 		}
 

management/server/http/events_handler_test.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -22,13 +23,13 @@ import (
 func initEventsTestData(account string, user *server.User, events ...*activity.Event) *EventsHandler {
 	return &EventsHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetEventsFunc: func(accountID, userID string) ([]*activity.Event, error) {
+			GetEventsFunc: func(_ context.Context, accountID, userID string) ([]*activity.Event, error) {
 				if accountID == account {
 					return events, nil
 				}
 				return []*activity.Event{}, nil
 			},
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return &server.Account{
 					Id:     claims.AccountId,
 					Domain: "hotmail.com",
@@ -37,7 +38,7 @@ func initEventsTestData(account string, user *server.User, events ...*activity.E
 					},
 				}, user, nil
 			},
-			GetUsersFromAccountFunc: func(accountID, userID string) ([]*server.UserInfo, error) {
+			GetUsersFromAccountFunc: func(_ context.Context, accountID, userID string) ([]*server.UserInfo, error) {
 				return make([]*server.UserInfo, 0), nil
 			},
 		},

management/server/http/geolocation_handler_test.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -35,13 +36,13 @@ func initGeolocationTestData(t *testing.T) *GeolocationsHandler {
 	err = util.CopyFileContents(geonamesDBPath, path.Join(tempDir, geolocation.GeoSqliteDBFile))
 	assert.NoError(t, err)
 
-	geo, err := geolocation.NewGeolocation(tempDir)
+	geo, err := geolocation.NewGeolocation(context.Background(), tempDir)
 	assert.NoError(t, err)
 	t.Cleanup(func() { _ = geo.Stop() })
 
 	return &GeolocationsHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				user := server.NewAdminUser("test_user")
 				return &server.Account{
 					Id: claims.AccountId,

management/server/http/geolocations_handler.go

@@ -40,19 +40,19 @@ func NewGeolocationsHandlerHandler(accountManager server.AccountManager, geoloca
 // GetAllCountries retrieves a list of all countries
 func (l *GeolocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Request) {
 	if err := l.authenticateUser(r); err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	if l.geolocationManager == nil {
 		// TODO: update error message to include geo db self hosted doc link when ready
-		util.WriteError(status.Errorf(status.PreconditionFailed, "Geo location database is not initialized"), w)
+		util.WriteError(r.Context(), status.Errorf(status.PreconditionFailed, "Geo location database is not initialized"), w)
 		return
 	}
 
 	allCountries, err := l.geolocationManager.GetAllCountries()
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -60,32 +60,32 @@ func (l *GeolocationsHandler) GetAllCountries(w http.ResponseWriter, r *http.Req
 	for _, country := range allCountries {
 		countries = append(countries, toCountryResponse(country))
 	}
-	util.WriteJSONObject(w, countries)
+	util.WriteJSONObject(r.Context(), w, countries)
 }
 
 // GetCitiesByCountry retrieves a list of cities based on the given country code
 func (l *GeolocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.Request) {
 	if err := l.authenticateUser(r); err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	countryCode := vars["country"]
 	if !countryCodeRegex.MatchString(countryCode) {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid country code"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid country code"), w)
 		return
 	}
 
 	if l.geolocationManager == nil {
-		util.WriteError(status.Errorf(status.PreconditionFailed, "Geo location database is not initialized. "+
+		util.WriteError(r.Context(), status.Errorf(status.PreconditionFailed, "Geo location database is not initialized. "+
 			"Check the self-hosted Geo database documentation at https://docs.netbird.io/selfhosted/geo-support"), w)
 		return
 	}
 
 	allCities, err := l.geolocationManager.GetCitiesByCountry(countryCode)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -93,12 +93,12 @@ func (l *GeolocationsHandler) GetCitiesByCountry(w http.ResponseWriter, r *http.
 	for _, city := range allCities {
 		cities = append(cities, toCityResponse(city))
 	}
-	util.WriteJSONObject(w, cities)
+	util.WriteJSONObject(r.Context(), w, cities)
 }
 
 func (l *GeolocationsHandler) authenticateUser(r *http.Request) error {
 	claims := l.claimsExtractor.FromRequestContext(r)
-	_, user, err := l.accountManager.GetAccountFromToken(claims)
+	_, user, err := l.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
 		return err
 	}

management/server/http/groups_handler.go

@@ -35,16 +35,16 @@ func NewGroupsHandler(accountManager server.AccountManager, authCfg AuthCfg) *Gr
 // GetAllGroups list for the account
 func (h *GroupsHandler) GetAllGroups(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		log.Error(err)
+		log.WithContext(r.Context()).Error(err)
 		http.Redirect(w, r, "/", http.StatusInternalServerError)
 		return
 	}
 
-	groups, err := h.accountManager.GetAllGroups(account.Id, user.Id)
+	groups, err := h.accountManager.GetAllGroups(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -53,42 +53,42 @@ func (h *GroupsHandler) GetAllGroups(w http.ResponseWriter, r *http.Request) {
 		groupsResponse = append(groupsResponse, toGroupResponse(account, group))
 	}
 
-	util.WriteJSONObject(w, groupsResponse)
+	util.WriteJSONObject(r.Context(), w, groupsResponse)
 }
 
 // UpdateGroup handles update to a group identified by a given ID
 func (h *GroupsHandler) UpdateGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	groupID, ok := vars["groupId"]
 	if !ok {
-		util.WriteError(status.Errorf(status.InvalidArgument, "group ID field is missing"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "group ID field is missing"), w)
 		return
 	}
 	if len(groupID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "group ID can't be empty"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "group ID can't be empty"), w)
 		return
 	}
 
 	eg, ok := account.Groups[groupID]
 	if !ok {
-		util.WriteError(status.Errorf(status.NotFound, "couldn't find group with ID %s", groupID), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "couldn't find group with ID %s", groupID), w)
 		return
 	}
 
 	allGroup, err := account.GetGroupAll()
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 	if allGroup.ID == groupID {
-		util.WriteError(status.Errorf(status.InvalidArgument, "updating group ALL is not allowed"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "updating group ALL is not allowed"), w)
 		return
 	}
 
@@ -100,7 +100,7 @@ func (h *GroupsHandler) UpdateGroup(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if req.Name == "" {
-		util.WriteError(status.Errorf(status.InvalidArgument, "group name shouldn't be empty"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "group name shouldn't be empty"), w)
 		return
 	}
 
@@ -118,21 +118,21 @@ func (h *GroupsHandler) UpdateGroup(w http.ResponseWriter, r *http.Request) {
 		IntegrationReference: eg.IntegrationReference,
 	}
 
-	if err := h.accountManager.SaveGroup(account.Id, user.Id, &group); err != nil {
-		log.Errorf("failed updating group %s under account %s %v", groupID, account.Id, err)
-		util.WriteError(err, w)
+	if err := h.accountManager.SaveGroup(r.Context(), account.Id, user.Id, &group); err != nil {
+		log.WithContext(r.Context()).Errorf("failed updating group %s under account %s %v", groupID, account.Id, err)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, toGroupResponse(account, &group))
+	util.WriteJSONObject(r.Context(), w, toGroupResponse(account, &group))
 }
 
 // CreateGroup handles group creation request
 func (h *GroupsHandler) CreateGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -144,7 +144,7 @@ func (h *GroupsHandler) CreateGroup(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if req.Name == "" {
-		util.WriteError(status.Errorf(status.InvalidArgument, "group name shouldn't be empty"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "group name shouldn't be empty"), w)
 		return
 	}
 
@@ -160,62 +160,62 @@ func (h *GroupsHandler) CreateGroup(w http.ResponseWriter, r *http.Request) {
 		Issued: nbgroup.GroupIssuedAPI,
 	}
 
-	err = h.accountManager.SaveGroup(account.Id, user.Id, &group)
+	err = h.accountManager.SaveGroup(r.Context(), account.Id, user.Id, &group)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, toGroupResponse(account, &group))
+	util.WriteJSONObject(r.Context(), w, toGroupResponse(account, &group))
 }
 
 // DeleteGroup handles group deletion request
 func (h *GroupsHandler) DeleteGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 	aID := account.Id
 
 	groupID := mux.Vars(r)["groupId"]
 	if len(groupID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid group ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid group ID"), w)
 		return
 	}
 
 	allGroup, err := account.GetGroupAll()
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	if allGroup.ID == groupID {
-		util.WriteError(status.Errorf(status.InvalidArgument, "deleting group ALL is not allowed"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "deleting group ALL is not allowed"), w)
 		return
 	}
 
-	err = h.accountManager.DeleteGroup(aID, user.Id, groupID)
+	err = h.accountManager.DeleteGroup(r.Context(), aID, user.Id, groupID)
 	if err != nil {
 		_, ok := err.(*server.GroupLinkError)
 		if ok {
 			util.WriteErrorResponse(err.Error(), http.StatusBadRequest, w)
 			return
 		}
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 // GetGroup returns a group
 func (h *GroupsHandler) GetGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -223,19 +223,19 @@ func (h *GroupsHandler) GetGroup(w http.ResponseWriter, r *http.Request) {
 	case http.MethodGet:
 		groupID := mux.Vars(r)["groupId"]
 		if len(groupID) == 0 {
-			util.WriteError(status.Errorf(status.InvalidArgument, "invalid group ID"), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid group ID"), w)
 			return
 		}
 
-		group, err := h.accountManager.GetGroup(account.Id, groupID, user.Id)
+		group, err := h.accountManager.GetGroup(r.Context(), account.Id, groupID, user.Id)
 		if err != nil {
-			util.WriteError(err, w)
+			util.WriteError(r.Context(), err, w)
 			return
 		}
 
-		util.WriteJSONObject(w, toGroupResponse(account, group))
+		util.WriteJSONObject(r.Context(), w, toGroupResponse(account, group))
 	default:
-		util.WriteError(status.Errorf(status.NotFound, "HTTP method not found"), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "HTTP method not found"), w)
 		return
 	}
 }

management/server/http/groups_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -32,13 +33,13 @@ var TestPeers = map[string]*nbpeer.Peer{
 func initGroupTestData(user *server.User, _ ...*nbgroup.Group) *GroupsHandler {
 	return &GroupsHandler{
 		accountManager: &mock_server.MockAccountManager{
-			SaveGroupFunc: func(accountID, userID string, group *nbgroup.Group) error {
+			SaveGroupFunc: func(_ context.Context, accountID, userID string, group *nbgroup.Group) error {
 				if !strings.HasPrefix(group.ID, "id-") {
 					group.ID = "id-was-set"
 				}
 				return nil
 			},
-			GetGroupFunc: func(_, groupID, _ string) (*nbgroup.Group, error) {
+			GetGroupFunc: func(_ context.Context, _, groupID, _ string) (*nbgroup.Group, error) {
 				if groupID != "idofthegroup" {
 					return nil, status.Errorf(status.NotFound, "not found")
 				}
@@ -55,7 +56,7 @@ func initGroupTestData(user *server.User, _ ...*nbgroup.Group) *GroupsHandler {
 					Issued: nbgroup.GroupIssuedAPI,
 				}, nil
 			},
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return &server.Account{
 					Id:     claims.AccountId,
 					Domain: "hotmail.com",
@@ -70,7 +71,7 @@ func initGroupTestData(user *server.User, _ ...*nbgroup.Group) *GroupsHandler {
 					},
 				}, user, nil
 			},
-			DeleteGroupFunc: func(accountID, userId, groupID string) error {
+			DeleteGroupFunc: func(_ context.Context, accountID, userId, groupID string) error {
 				if groupID == "linked-grp" {
 					return &server.GroupLinkError{
 						Resource: "something",

management/server/http/handler.go

@@ -9,6 +9,7 @@ import (
 	"github.com/rs/cors"
 
 	"github.com/netbirdio/management-integrations/integrations"
+
 	s "github.com/netbirdio/netbird/management/server"
 	"github.com/netbirdio/netbird/management/server/geolocation"
 	"github.com/netbirdio/netbird/management/server/http/middleware"
@@ -57,6 +58,11 @@ func APIHandler(ctx context.Context, accountManager s.AccountManager, LocationMa
 
 	corsMiddleware := cors.AllowAll()
 
+	claimsExtractor = jwtclaims.NewClaimsExtractor(
+		jwtclaims.WithAudience(authCfg.Audience),
+		jwtclaims.WithUserIDClaim(authCfg.UserIDClaim),
+	)
+
 	acMiddleware := middleware.NewAccessControl(
 		authCfg.Audience,
 		authCfg.UserIDClaim,

management/server/http/middleware/access_control.go

@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"context"
 	"net/http"
 	"regexp"
 
@@ -15,7 +16,7 @@ import (
 )
 
 // GetUser function defines a function to fetch user from Account by jwtclaims.AuthorizationClaims
-type GetUser func(claims jwtclaims.AuthorizationClaims) (*server.User, error)
+type GetUser func(ctx context.Context, claims jwtclaims.AuthorizationClaims) (*server.User, error)
 
 // AccessControl middleware to restrict to make POST/PUT/DELETE requests by admin only
 type AccessControl struct {
@@ -46,15 +47,15 @@ func (a *AccessControl) Handler(h http.Handler) http.Handler {
 
 		claims := a.claimsExtract.FromRequestContext(r)
 
-		user, err := a.getUser(claims)
+		user, err := a.getUser(r.Context(), claims)
 		if err != nil {
-			log.Errorf("failed to get user from claims: %s", err)
-			util.WriteError(status.Errorf(status.Unauthorized, "invalid JWT"), w)
+			log.WithContext(r.Context()).Errorf("failed to get user from claims: %s", err)
+			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "invalid JWT"), w)
 			return
 		}
 
 		if user.IsBlocked() {
-			util.WriteError(status.Errorf(status.PermissionDenied, "the user has no access to the API or is blocked"), w)
+			util.WriteError(r.Context(), status.Errorf(status.PermissionDenied, "the user has no access to the API or is blocked"), w)
 			return
 		}
 
@@ -63,12 +64,12 @@ func (a *AccessControl) Handler(h http.Handler) http.Handler {
 			case http.MethodDelete, http.MethodPost, http.MethodPatch, http.MethodPut:
 
 				if tokenPathRegexp.MatchString(r.URL.Path) {
-					log.Debugf("valid Path")
+					log.WithContext(r.Context()).Debugf("valid Path")
 					h.ServeHTTP(w, r)
 					return
 				}
 
-				util.WriteError(status.Errorf(status.PermissionDenied, "only users with admin power can perform this operation"), w)
+				util.WriteError(r.Context(), status.Errorf(status.PermissionDenied, "only users with admin power can perform this operation"), w)
 				return
 			}
 		}

management/server/http/middleware/auth_middleware.go

@@ -12,6 +12,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server"
+	nbContext "github.com/netbirdio/netbird/management/server/context"
 	"github.com/netbirdio/netbird/management/server/http/middleware/bypass"
 	"github.com/netbirdio/netbird/management/server/http/util"
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
@@ -19,16 +20,16 @@ import (
 )
 
 // GetAccountFromPATFunc function
-type GetAccountFromPATFunc func(token string) (*server.Account, *server.User, *server.PersonalAccessToken, error)
+type GetAccountFromPATFunc func(ctx context.Context, token string) (*server.Account, *server.User, *server.PersonalAccessToken, error)
 
 // ValidateAndParseTokenFunc function
-type ValidateAndParseTokenFunc func(token string) (*jwt.Token, error)
+type ValidateAndParseTokenFunc func(ctx context.Context, token string) (*jwt.Token, error)
 
 // MarkPATUsedFunc function
-type MarkPATUsedFunc func(token string) error
+type MarkPATUsedFunc func(ctx context.Context, token string) error
 
 // CheckUserAccessByJWTGroupsFunc function
-type CheckUserAccessByJWTGroupsFunc func(claims jwtclaims.AuthorizationClaims) error
+type CheckUserAccessByJWTGroupsFunc func(ctx context.Context, claims jwtclaims.AuthorizationClaims) error
 
 // AuthMiddleware middleware to verify personal access tokens (PAT) and JWT tokens
 type AuthMiddleware struct {
@@ -85,23 +86,27 @@ func (m *AuthMiddleware) Handler(h http.Handler) http.Handler {
 		case "bearer":
 			err := m.checkJWTFromRequest(w, r, auth)
 			if err != nil {
-				log.Errorf("Error when validating JWT claims: %s", err.Error())
-				util.WriteError(status.Errorf(status.Unauthorized, "token invalid"), w)
+				log.WithContext(r.Context()).Errorf("Error when validating JWT claims: %s", err.Error())
+				util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "token invalid"), w)
 				return
 			}
-			h.ServeHTTP(w, r)
 		case "token":
 			err := m.checkPATFromRequest(w, r, auth)
 			if err != nil {
-				log.Debugf("Error when validating PAT claims: %s", err.Error())
-				util.WriteError(status.Errorf(status.Unauthorized, "token invalid"), w)
+				log.WithContext(r.Context()).Debugf("Error when validating PAT claims: %s", err.Error())
+				util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "token invalid"), w)
 				return
 			}
-			h.ServeHTTP(w, r)
 		default:
-			util.WriteError(status.Errorf(status.Unauthorized, "no valid authentication provided"), w)
+			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "no valid authentication provided"), w)
 			return
 		}
+		claims := m.claimsExtractor.FromRequestContext(r)
+		//nolint
+		ctx := context.WithValue(r.Context(), nbContext.UserIDKey, claims.UserId)
+		//nolint
+		ctx = context.WithValue(ctx, nbContext.AccountIDKey, claims.AccountId)
+		h.ServeHTTP(w, r.WithContext(ctx))
 	})
 }
 
@@ -114,7 +119,7 @@ func (m *AuthMiddleware) checkJWTFromRequest(w http.ResponseWriter, r *http.Requ
 		return fmt.Errorf("Error extracting token: %w", err)
 	}
 
-	validatedToken, err := m.validateAndParseToken(token)
+	validatedToken, err := m.validateAndParseToken(r.Context(), token)
 	if err != nil {
 		return err
 	}
@@ -123,7 +128,7 @@ func (m *AuthMiddleware) checkJWTFromRequest(w http.ResponseWriter, r *http.Requ
 		return nil
 	}
 
-	if err := m.verifyUserAccess(validatedToken); err != nil {
+	if err := m.verifyUserAccess(r.Context(), validatedToken); err != nil {
 		return err
 	}
 
@@ -138,9 +143,9 @@ func (m *AuthMiddleware) checkJWTFromRequest(w http.ResponseWriter, r *http.Requ
 // verifyUserAccess checks if a user, based on a validated JWT token,
 // is allowed access, particularly in cases where the admin enabled JWT
 // group propagation and designated certain groups with access permissions.
-func (m *AuthMiddleware) verifyUserAccess(validatedToken *jwt.Token) error {
+func (m *AuthMiddleware) verifyUserAccess(ctx context.Context, validatedToken *jwt.Token) error {
 	authClaims := m.claimsExtractor.FromToken(validatedToken)
-	return m.checkUserAccessByJWTGroups(authClaims)
+	return m.checkUserAccessByJWTGroups(ctx, authClaims)
 }
 
 // CheckPATFromRequest checks if the PAT is valid
@@ -152,7 +157,7 @@ func (m *AuthMiddleware) checkPATFromRequest(w http.ResponseWriter, r *http.Requ
 		return fmt.Errorf("Error extracting token: %w", err)
 	}
 
-	account, user, pat, err := m.getAccountFromPAT(token)
+	account, user, pat, err := m.getAccountFromPAT(r.Context(), token)
 	if err != nil {
 		return fmt.Errorf("invalid Token: %w", err)
 	}
@@ -160,7 +165,7 @@ func (m *AuthMiddleware) checkPATFromRequest(w http.ResponseWriter, r *http.Requ
 		return fmt.Errorf("token expired")
 	}
 
-	err = m.markPATUsed(pat.ID)
+	err = m.markPATUsed(r.Context(), pat.ID)
 	if err != nil {
 		return err
 	}

management/server/http/middleware/auth_middleware_test.go

@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -15,15 +16,16 @@ import (
 )
 
 const (
-	audience    = "audience"
-	userIDClaim = "userIDClaim"
-	accountID   = "accountID"
-	domain      = "domain"
-	userID      = "userID"
-	tokenID     = "tokenID"
-	PAT         = "nbp_PAT"
-	JWT         = "JWT"
-	wrongToken  = "wrongToken"
+	audience       = "audience"
+	userIDClaim    = "userIDClaim"
+	accountID      = "accountID"
+	domain         = "domain"
+	domainCategory = "domainCategory"
+	userID         = "userID"
+	tokenID        = "tokenID"
+	PAT            = "nbp_PAT"
+	JWT            = "JWT"
+	wrongToken     = "wrongToken"
 )
 
 var testAccount = &server.Account{
@@ -47,14 +49,14 @@ var testAccount = &server.Account{
 	},
 }
 
-func mockGetAccountFromPAT(token string) (*server.Account, *server.User, *server.PersonalAccessToken, error) {
+func mockGetAccountFromPAT(_ context.Context, token string) (*server.Account, *server.User, *server.PersonalAccessToken, error) {
 	if token == PAT {
 		return testAccount, testAccount.Users[userID], testAccount.Users[userID].PATs[tokenID], nil
 	}
 	return nil, nil, nil, fmt.Errorf("PAT invalid")
 }
 
-func mockValidateAndParseToken(token string) (*jwt.Token, error) {
+func mockValidateAndParseToken(_ context.Context, token string) (*jwt.Token, error) {
 	if token == JWT {
 		return &jwt.Token{
 			Claims: jwt.MapClaims{
@@ -67,14 +69,14 @@ func mockValidateAndParseToken(token string) (*jwt.Token, error) {
 	return nil, fmt.Errorf("JWT invalid")
 }
 
-func mockMarkPATUsed(token string) error {
+func mockMarkPATUsed(_ context.Context, token string) error {
 	if token == tokenID {
 		return nil
 	}
 	return fmt.Errorf("Should never get reached")
 }
 
-func mockCheckUserAccessByJWTGroups(claims jwtclaims.AuthorizationClaims) error {
+func mockCheckUserAccessByJWTGroups(_ context.Context, claims jwtclaims.AuthorizationClaims) error {
 	if testAccount.Id != claims.AccountId {
 		return fmt.Errorf("account with id %s does not exist", claims.AccountId)
 	}

management/server/http/middleware/bypass/bypass.go

@@ -56,7 +56,7 @@ func ShouldBypass(requestPath string, h http.Handler, w http.ResponseWriter, r *
 	for bypassPath := range bypassPaths {
 		matched, err := path.Match(bypassPath, requestPath)
 		if err != nil {
-			log.Errorf("Error matching path %s with %s from %s: %v", bypassPath, requestPath, GetList(), err)
+			log.WithContext(r.Context()).Errorf("Error matching path %s with %s from %s: %v", bypassPath, requestPath, GetList(), err)
 			continue
 		}
 		if matched {

management/server/http/nameservers_handler.go

@@ -36,16 +36,16 @@ func NewNameserversHandler(accountManager server.AccountManager, authCfg AuthCfg
 // GetAllNameservers returns the list of nameserver groups for the account
 func (h *NameserversHandler) GetAllNameservers(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		log.Error(err)
+		log.WithContext(r.Context()).Error(err)
 		http.Redirect(w, r, "/", http.StatusInternalServerError)
 		return
 	}
 
-	nsGroups, err := h.accountManager.ListNameServerGroups(account.Id, user.Id)
+	nsGroups, err := h.accountManager.ListNameServerGroups(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -54,15 +54,15 @@ func (h *NameserversHandler) GetAllNameservers(w http.ResponseWriter, r *http.Re
 		apiNameservers = append(apiNameservers, toNameserverGroupResponse(r))
 	}
 
-	util.WriteJSONObject(w, apiNameservers)
+	util.WriteJSONObject(r.Context(), w, apiNameservers)
 }
 
 // CreateNameserverGroup handles nameserver group creation request
 func (h *NameserversHandler) CreateNameserverGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -75,33 +75,33 @@ func (h *NameserversHandler) CreateNameserverGroup(w http.ResponseWriter, r *htt
 
 	nsList, err := toServerNSList(req.Nameservers)
 	if err != nil {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid NS servers format"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid NS servers format"), w)
 		return
 	}
 
-	nsGroup, err := h.accountManager.CreateNameServerGroup(account.Id, req.Name, req.Description, nsList, req.Groups, req.Primary, req.Domains, req.Enabled, user.Id, req.SearchDomainsEnabled)
+	nsGroup, err := h.accountManager.CreateNameServerGroup(r.Context(), account.Id, req.Name, req.Description, nsList, req.Groups, req.Primary, req.Domains, req.Enabled, user.Id, req.SearchDomainsEnabled)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	resp := toNameserverGroupResponse(nsGroup)
 
-	util.WriteJSONObject(w, &resp)
+	util.WriteJSONObject(r.Context(), w, &resp)
 }
 
 // UpdateNameserverGroup handles update to a nameserver group identified by a given ID
 func (h *NameserversHandler) UpdateNameserverGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	nsGroupID := mux.Vars(r)["nsgroupId"]
 	if len(nsGroupID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid nameserver group ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid nameserver group ID"), w)
 		return
 	}
 
@@ -114,7 +114,7 @@ func (h *NameserversHandler) UpdateNameserverGroup(w http.ResponseWriter, r *htt
 
 	nsList, err := toServerNSList(req.Nameservers)
 	if err != nil {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid NS servers format"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid NS servers format"), w)
 		return
 	}
 
@@ -130,66 +130,66 @@ func (h *NameserversHandler) UpdateNameserverGroup(w http.ResponseWriter, r *htt
 		SearchDomainsEnabled: req.SearchDomainsEnabled,
 	}
 
-	err = h.accountManager.SaveNameServerGroup(account.Id, user.Id, updatedNSGroup)
+	err = h.accountManager.SaveNameServerGroup(r.Context(), account.Id, user.Id, updatedNSGroup)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	resp := toNameserverGroupResponse(updatedNSGroup)
 
-	util.WriteJSONObject(w, &resp)
+	util.WriteJSONObject(r.Context(), w, &resp)
 }
 
 // DeleteNameserverGroup handles nameserver group deletion request
 func (h *NameserversHandler) DeleteNameserverGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	nsGroupID := mux.Vars(r)["nsgroupId"]
 	if len(nsGroupID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid nameserver group ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid nameserver group ID"), w)
 		return
 	}
 
-	err = h.accountManager.DeleteNameServerGroup(account.Id, nsGroupID, user.Id)
+	err = h.accountManager.DeleteNameServerGroup(r.Context(), account.Id, nsGroupID, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 // GetNameserverGroup handles a nameserver group Get request identified by ID
 func (h *NameserversHandler) GetNameserverGroup(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		log.Error(err)
+		log.WithContext(r.Context()).Error(err)
 		http.Redirect(w, r, "/", http.StatusInternalServerError)
 		return
 	}
 
 	nsGroupID := mux.Vars(r)["nsgroupId"]
 	if len(nsGroupID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid nameserver group ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid nameserver group ID"), w)
 		return
 	}
 
-	nsGroup, err := h.accountManager.GetNameServerGroup(account.Id, user.Id, nsGroupID)
+	nsGroup, err := h.accountManager.GetNameServerGroup(r.Context(), account.Id, user.Id, nsGroupID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	resp := toNameserverGroupResponse(nsGroup)
 
-	util.WriteJSONObject(w, &resp)
+	util.WriteJSONObject(r.Context(), w, &resp)
 }
 
 func toServerNSList(apiNSList []api.Nameserver) ([]nbdns.NameServer, error) {

management/server/http/nameservers_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -61,13 +62,13 @@ var baseExistingNSGroup = &nbdns.NameServerGroup{
 func initNameserversTestData() *NameserversHandler {
 	return &NameserversHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetNameServerGroupFunc: func(accountID, userID, nsGroupID string) (*nbdns.NameServerGroup, error) {
+			GetNameServerGroupFunc: func(_ context.Context, accountID, userID, nsGroupID string) (*nbdns.NameServerGroup, error) {
 				if nsGroupID == existingNSGroupID {
 					return baseExistingNSGroup.Copy(), nil
 				}
 				return nil, status.Errorf(status.NotFound, "nameserver group with ID %s not found", nsGroupID)
 			},
-			CreateNameServerGroupFunc: func(accountID string, name, description string, nameServerList []nbdns.NameServer, groups []string, primary bool, domains []string, enabled bool, _ string, searchDomains bool) (*nbdns.NameServerGroup, error) {
+			CreateNameServerGroupFunc: func(_ context.Context, accountID string, name, description string, nameServerList []nbdns.NameServer, groups []string, primary bool, domains []string, enabled bool, _ string, searchDomains bool) (*nbdns.NameServerGroup, error) {
 				return &nbdns.NameServerGroup{
 					ID:                   existingNSGroupID,
 					Name:                 name,
@@ -80,16 +81,16 @@ func initNameserversTestData() *NameserversHandler {
 					SearchDomainsEnabled: searchDomains,
 				}, nil
 			},
-			DeleteNameServerGroupFunc: func(accountID, nsGroupID, _ string) error {
+			DeleteNameServerGroupFunc: func(_ context.Context, accountID, nsGroupID, _ string) error {
 				return nil
 			},
-			SaveNameServerGroupFunc: func(accountID, _ string, nsGroupToSave *nbdns.NameServerGroup) error {
+			SaveNameServerGroupFunc: func(_ context.Context, accountID, _ string, nsGroupToSave *nbdns.NameServerGroup) error {
 				if nsGroupToSave.ID == existingNSGroupID {
 					return nil
 				}
 				return status.Errorf(status.NotFound, "nameserver group with ID %s was not found", nsGroupToSave.ID)
 			},
-			GetAccountFromTokenFunc: func(_ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, _ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return testingNSAccount, testingAccount.Users["test_user"], nil
 			},
 		},

management/server/http/pat_handler.go

@@ -34,22 +34,22 @@ func NewPATsHandler(accountManager server.AccountManager, authCfg AuthCfg) *PATH
 // GetAllTokens is HTTP GET handler that returns a list of all personal access tokens for the given user
 func (h *PATHandler) GetAllTokens(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	userID := vars["userId"]
 	if len(userID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
-	pats, err := h.accountManager.GetAllPATs(account.Id, user.Id, userID)
+	pats, err := h.accountManager.GetAllPATs(r.Context(), account.Id, user.Id, userID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -58,53 +58,53 @@ func (h *PATHandler) GetAllTokens(w http.ResponseWriter, r *http.Request) {
 		patResponse = append(patResponse, toPATResponse(pat))
 	}
 
-	util.WriteJSONObject(w, patResponse)
+	util.WriteJSONObject(r.Context(), w, patResponse)
 }
 
 // GetToken is HTTP GET handler that returns a personal access token for the given user
 func (h *PATHandler) GetToken(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	targetUserID := vars["userId"]
 	if len(targetUserID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
 	tokenID := vars["tokenId"]
 	if len(tokenID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid token ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid token ID"), w)
 		return
 	}
 
-	pat, err := h.accountManager.GetPAT(account.Id, user.Id, targetUserID, tokenID)
+	pat, err := h.accountManager.GetPAT(r.Context(), account.Id, user.Id, targetUserID, tokenID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, toPATResponse(pat))
+	util.WriteJSONObject(r.Context(), w, toPATResponse(pat))
 }
 
 // CreateToken is HTTP POST handler that creates a personal access token for the given user
 func (h *PATHandler) CreateToken(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	targetUserID := vars["userId"]
 	if len(targetUserID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
@@ -115,44 +115,44 @@ func (h *PATHandler) CreateToken(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	pat, err := h.accountManager.CreatePAT(account.Id, user.Id, targetUserID, req.Name, req.ExpiresIn)
+	pat, err := h.accountManager.CreatePAT(r.Context(), account.Id, user.Id, targetUserID, req.Name, req.ExpiresIn)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, toPATGeneratedResponse(pat))
+	util.WriteJSONObject(r.Context(), w, toPATGeneratedResponse(pat))
 }
 
 // DeleteToken is HTTP DELETE handler that deletes a personal access token for the given user
 func (h *PATHandler) DeleteToken(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	targetUserID := vars["userId"]
 	if len(targetUserID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
 	tokenID := vars["tokenId"]
 	if len(tokenID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid token ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid token ID"), w)
 		return
 	}
 
-	err = h.accountManager.DeletePAT(account.Id, user.Id, targetUserID, tokenID)
+	err = h.accountManager.DeletePAT(r.Context(), account.Id, user.Id, targetUserID, tokenID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 func toPATResponse(pat *server.PersonalAccessToken) *api.PersonalAccessToken {

management/server/http/pat_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -63,7 +64,7 @@ var testAccount = &server.Account{
 func initPATTestData() *PATHandler {
 	return &PATHandler{
 		accountManager: &mock_server.MockAccountManager{
-			CreatePATFunc: func(accountID string, initiatorUserID string, targetUserID string, tokenName string, expiresIn int) (*server.PersonalAccessTokenGenerated, error) {
+			CreatePATFunc: func(_ context.Context, accountID string, initiatorUserID string, targetUserID string, tokenName string, expiresIn int) (*server.PersonalAccessTokenGenerated, error) {
 				if accountID != existingAccountID {
 					return nil, status.Errorf(status.NotFound, "account with ID %s not found", accountID)
 				}
@@ -76,10 +77,10 @@ func initPATTestData() *PATHandler {
 				}, nil
 			},
 
-			GetAccountFromTokenFunc: func(_ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, _ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return testAccount, testAccount.Users[existingUserID], nil
 			},
-			DeletePATFunc: func(accountID string, initiatorUserID string, targetUserID string, tokenID string) error {
+			DeletePATFunc: func(_ context.Context, accountID string, initiatorUserID string, targetUserID string, tokenID string) error {
 				if accountID != existingAccountID {
 					return status.Errorf(status.NotFound, "account with ID %s not found", accountID)
 				}
@@ -91,7 +92,7 @@ func initPATTestData() *PATHandler {
 				}
 				return nil
 			},
-			GetPATFunc: func(accountID string, initiatorUserID string, targetUserID string, tokenID string) (*server.PersonalAccessToken, error) {
+			GetPATFunc: func(_ context.Context, accountID string, initiatorUserID string, targetUserID string, tokenID string) (*server.PersonalAccessToken, error) {
 				if accountID != existingAccountID {
 					return nil, status.Errorf(status.NotFound, "account with ID %s not found", accountID)
 				}
@@ -103,7 +104,7 @@ func initPATTestData() *PATHandler {
 				}
 				return testAccount.Users[existingUserID].PATs[existingTokenID], nil
 			},
-			GetAllPATsFunc: func(accountID string, initiatorUserID string, targetUserID string) ([]*server.PersonalAccessToken, error) {
+			GetAllPATsFunc: func(_ context.Context, accountID string, initiatorUserID string, targetUserID string) ([]*server.PersonalAccessToken, error) {
 				if accountID != existingAccountID {
 					return nil, status.Errorf(status.NotFound, "account with ID %s not found", accountID)
 				}

management/server/http/peers_handler.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -47,16 +48,16 @@ func (h *PeersHandler) checkPeerStatus(peer *nbpeer.Peer) (*nbpeer.Peer, error)
 	return peerToReturn, nil
 }
 
-func (h *PeersHandler) getPeer(account *server.Account, peerID, userID string, w http.ResponseWriter) {
-	peer, err := h.accountManager.GetPeer(account.Id, peerID, userID)
+func (h *PeersHandler) getPeer(ctx context.Context, account *server.Account, peerID, userID string, w http.ResponseWriter) {
+	peer, err := h.accountManager.GetPeer(ctx, account.Id, peerID, userID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(ctx, err, w)
 		return
 	}
 
 	peerToReturn, err := h.checkPeerStatus(peer)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(ctx, err, w)
 		return
 	}
 	dnsDomain := h.accountManager.GetDNSDomain()
@@ -65,19 +66,19 @@ func (h *PeersHandler) getPeer(account *server.Account, peerID, userID string, w
 
 	validPeers, err := h.accountManager.GetValidatedPeers(account)
 	if err != nil {
-		log.Errorf("failed to list appreoved peers: %v", err)
-		util.WriteError(fmt.Errorf("internal error"), w)
+		log.WithContext(ctx).Errorf("failed to list appreoved peers: %v", err)
+		util.WriteError(ctx, fmt.Errorf("internal error"), w)
 		return
 	}
 
-	netMap := account.GetPeerNetworkMap(peerID, h.accountManager.GetDNSDomain(), validPeers)
+	netMap := account.GetPeerNetworkMap(ctx, peerID, h.accountManager.GetDNSDomain(), validPeers)
 	accessiblePeers := toAccessiblePeers(netMap, dnsDomain)
 
 	_, valid := validPeers[peer.ID]
-	util.WriteJSONObject(w, toSinglePeerResponse(peerToReturn, groupsInfo, dnsDomain, accessiblePeers, valid))
+	util.WriteJSONObject(ctx, w, toSinglePeerResponse(peerToReturn, groupsInfo, dnsDomain, accessiblePeers, valid))
 }
 
-func (h *PeersHandler) updatePeer(account *server.Account, user *server.User, peerID string, w http.ResponseWriter, r *http.Request) {
+func (h *PeersHandler) updatePeer(ctx context.Context, account *server.Account, user *server.User, peerID string, w http.ResponseWriter, r *http.Request) {
 	req := &api.PeerRequest{}
 	err := json.NewDecoder(r.Body).Decode(&req)
 	if err != nil {
@@ -99,9 +100,9 @@ func (h *PeersHandler) updatePeer(account *server.Account, user *server.User, pe
 		}
 	}
 
-	peer, err := h.accountManager.UpdatePeer(account.Id, user.Id, update)
+	peer, err := h.accountManager.UpdatePeer(ctx, account.Id, user.Id, update)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(ctx, err, w)
 		return
 	}
 	dnsDomain := h.accountManager.GetDNSDomain()
@@ -110,75 +111,75 @@ func (h *PeersHandler) updatePeer(account *server.Account, user *server.User, pe
 
 	validPeers, err := h.accountManager.GetValidatedPeers(account)
 	if err != nil {
-		log.Errorf("failed to list appreoved peers: %v", err)
-		util.WriteError(fmt.Errorf("internal error"), w)
+		log.WithContext(ctx).Errorf("failed to list appreoved peers: %v", err)
+		util.WriteError(ctx, fmt.Errorf("internal error"), w)
 		return
 	}
-	netMap := account.GetPeerNetworkMap(peerID, h.accountManager.GetDNSDomain(), validPeers)
+	netMap := account.GetPeerNetworkMap(ctx, peerID, h.accountManager.GetDNSDomain(), validPeers)
 	accessiblePeers := toAccessiblePeers(netMap, dnsDomain)
 
 	_, valid := validPeers[peer.ID]
 
-	util.WriteJSONObject(w, toSinglePeerResponse(peer, groupMinimumInfo, dnsDomain, accessiblePeers, valid))
+	util.WriteJSONObject(r.Context(), w, toSinglePeerResponse(peer, groupMinimumInfo, dnsDomain, accessiblePeers, valid))
 }
 
-func (h *PeersHandler) deletePeer(accountID, userID string, peerID string, w http.ResponseWriter) {
-	err := h.accountManager.DeletePeer(accountID, peerID, userID)
+func (h *PeersHandler) deletePeer(ctx context.Context, accountID, userID string, peerID string, w http.ResponseWriter) {
+	err := h.accountManager.DeletePeer(ctx, accountID, peerID, userID)
 	if err != nil {
-		log.Errorf("failed to delete peer: %v", err)
-		util.WriteError(err, w)
+		log.WithContext(ctx).Errorf("failed to delete peer: %v", err)
+		util.WriteError(ctx, err, w)
 		return
 	}
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(ctx, w, emptyObject{})
 }
 
 // HandlePeer handles all peer requests for GET, PUT and DELETE operations
 func (h *PeersHandler) HandlePeer(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 	vars := mux.Vars(r)
 	peerID := vars["peerId"]
 	if len(peerID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid peer ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid peer ID"), w)
 		return
 	}
 
 	switch r.Method {
 	case http.MethodDelete:
-		h.deletePeer(account.Id, user.Id, peerID, w)
+		h.deletePeer(r.Context(), account.Id, user.Id, peerID, w)
 		return
 	case http.MethodPut:
-		h.updatePeer(account, user, peerID, w, r)
+		h.updatePeer(r.Context(), account, user, peerID, w, r)
 		return
 	case http.MethodGet:
-		h.getPeer(account, peerID, user.Id, w)
+		h.getPeer(r.Context(), account, peerID, user.Id, w)
 		return
 	default:
-		util.WriteError(status.Errorf(status.NotFound, "unknown METHOD"), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "unknown METHOD"), w)
 	}
 }
 
 // GetAllPeers returns a list of all peers associated with a provided account
 func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodGet {
-		util.WriteError(status.Errorf(status.NotFound, "unknown METHOD"), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "unknown METHOD"), w)
 		return
 	}
 
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	peers, err := h.accountManager.GetPeers(account.Id, user.Id)
+	peers, err := h.accountManager.GetPeers(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -188,34 +189,34 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 	for _, peer := range peers {
 		peerToReturn, err := h.checkPeerStatus(peer)
 		if err != nil {
-			util.WriteError(err, w)
+			util.WriteError(r.Context(), err, w)
 			return
 		}
 		groupMinimumInfo := toGroupsInfo(account.Groups, peer.ID)
 
-		accessiblePeerNumbers, _ := h.accessiblePeersNumber(account, peer.ID)
+		accessiblePeerNumbers, _ := h.accessiblePeersNumber(r.Context(), account, peer.ID)
 
 		respBody = append(respBody, toPeerListItemResponse(peerToReturn, groupMinimumInfo, dnsDomain, accessiblePeerNumbers))
 	}
 
 	validPeersMap, err := h.accountManager.GetValidatedPeers(account)
 	if err != nil {
-		log.Errorf("failed to list appreoved peers: %v", err)
-		util.WriteError(fmt.Errorf("internal error"), w)
+		log.WithContext(r.Context()).Errorf("failed to list appreoved peers: %v", err)
+		util.WriteError(r.Context(), fmt.Errorf("internal error"), w)
 		return
 	}
 	h.setApprovalRequiredFlag(respBody, validPeersMap)
 
-	util.WriteJSONObject(w, respBody)
+	util.WriteJSONObject(r.Context(), w, respBody)
 }
 
-func (h *PeersHandler) accessiblePeersNumber(account *server.Account, peerID string) (int, error) {
+func (h *PeersHandler) accessiblePeersNumber(ctx context.Context, account *server.Account, peerID string) (int, error) {
 	validatedPeersMap, err := h.accountManager.GetValidatedPeers(account)
 	if err != nil {
 		return 0, err
 	}
 
-	netMap := account.GetPeerNetworkMap(peerID, h.accountManager.GetDNSDomain(), validatedPeersMap)
+	netMap := account.GetPeerNetworkMap(ctx, peerID, h.accountManager.GetDNSDomain(), validatedPeersMap)
 	return len(netMap.Peers) + len(netMap.OfflinePeers), nil
 }
 

management/server/http/peers_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net"
@@ -29,7 +30,7 @@ const noUpdateChannelTestPeerID = "no-update-channel"
 func initTestMetaData(peers ...*nbpeer.Peer) *PeersHandler {
 	return &PeersHandler{
 		accountManager: &mock_server.MockAccountManager{
-			UpdatePeerFunc: func(accountID, userID string, update *nbpeer.Peer) (*nbpeer.Peer, error) {
+			UpdatePeerFunc: func(_ context.Context, accountID, userID string, update *nbpeer.Peer) (*nbpeer.Peer, error) {
 				var p *nbpeer.Peer
 				for _, peer := range peers {
 					if update.ID == peer.ID {
@@ -42,7 +43,7 @@ func initTestMetaData(peers ...*nbpeer.Peer) *PeersHandler {
 				p.Name = update.Name
 				return p, nil
 			},
-			GetPeerFunc: func(accountID, peerID, userID string) (*nbpeer.Peer, error) {
+			GetPeerFunc: func(_ context.Context, accountID, peerID, userID string) (*nbpeer.Peer, error) {
 				var p *nbpeer.Peer
 				for _, peer := range peers {
 					if peerID == peer.ID {
@@ -52,13 +53,13 @@ func initTestMetaData(peers ...*nbpeer.Peer) *PeersHandler {
 				}
 				return p, nil
 			},
-			GetPeersFunc: func(accountID, userID string) ([]*nbpeer.Peer, error) {
+			GetPeersFunc: func(_ context.Context, accountID, userID string) ([]*nbpeer.Peer, error) {
 				return peers, nil
 			},
 			GetDNSDomainFunc: func() string {
 				return "netbird.selfhosted"
 			},
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				user := server.NewAdminUser("test_user")
 				return &server.Account{
 					Id:     claims.AccountId,

management/server/http/policies_handler.go

@@ -35,15 +35,15 @@ func NewPoliciesHandler(accountManager server.AccountManager, authCfg AuthCfg) *
 // GetAllPolicies list for the account
 func (h *Policies) GetAllPolicies(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	accountPolicies, err := h.accountManager.ListPolicies(account.Id, user.Id)
+	accountPolicies, err := h.accountManager.ListPolicies(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -51,28 +51,28 @@ func (h *Policies) GetAllPolicies(w http.ResponseWriter, r *http.Request) {
 	for _, policy := range accountPolicies {
 		resp := toPolicyResponse(account, policy)
 		if len(resp.Rules) == 0 {
-			util.WriteError(status.Errorf(status.Internal, "no rules in the policy"), w)
+			util.WriteError(r.Context(), status.Errorf(status.Internal, "no rules in the policy"), w)
 			return
 		}
 		policies = append(policies, resp)
 	}
 
-	util.WriteJSONObject(w, policies)
+	util.WriteJSONObject(r.Context(), w, policies)
 }
 
 // UpdatePolicy handles update to a policy identified by a given ID
 func (h *Policies) UpdatePolicy(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	policyID := vars["policyId"]
 	if len(policyID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid policy ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid policy ID"), w)
 		return
 	}
 
@@ -84,7 +84,7 @@ func (h *Policies) UpdatePolicy(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 	if policyIdx < 0 {
-		util.WriteError(status.Errorf(status.NotFound, "couldn't find policy id %s", policyID), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "couldn't find policy id %s", policyID), w)
 		return
 	}
 
@@ -94,9 +94,9 @@ func (h *Policies) UpdatePolicy(w http.ResponseWriter, r *http.Request) {
 // CreatePolicy handles policy creation request
 func (h *Policies) CreatePolicy(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -118,12 +118,12 @@ func (h *Policies) savePolicy(
 	}
 
 	if req.Name == "" {
-		util.WriteError(status.Errorf(status.InvalidArgument, "policy name shouldn't be empty"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "policy name shouldn't be empty"), w)
 		return
 	}
 
 	if len(req.Rules) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "policy rules shouldn't be empty"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "policy rules shouldn't be empty"), w)
 		return
 	}
 
@@ -137,31 +137,31 @@ func (h *Policies) savePolicy(
 		Enabled:     req.Enabled,
 		Description: req.Description,
 	}
-	for _, r := range req.Rules {
+	for _, rule := range req.Rules {
 		pr := server.PolicyRule{
-			ID:            policyID, //TODO: when policy can contain multiple rules, need refactor
-			Name:          r.Name,
-			Destinations:  groupMinimumsToStrings(account, r.Destinations),
-			Sources:       groupMinimumsToStrings(account, r.Sources),
-			Bidirectional: r.Bidirectional,
+			ID:            policyID, // TODO: when policy can contain multiple rules, need refactor
+			Name:          rule.Name,
+			Destinations:  groupMinimumsToStrings(account, rule.Destinations),
+			Sources:       groupMinimumsToStrings(account, rule.Sources),
+			Bidirectional: rule.Bidirectional,
 		}
 
-		pr.Enabled = r.Enabled
-		if r.Description != nil {
-			pr.Description = *r.Description
+		pr.Enabled = rule.Enabled
+		if rule.Description != nil {
+			pr.Description = *rule.Description
 		}
 
-		switch r.Action {
+		switch rule.Action {
 		case api.PolicyRuleUpdateActionAccept:
 			pr.Action = server.PolicyTrafficActionAccept
 		case api.PolicyRuleUpdateActionDrop:
 			pr.Action = server.PolicyTrafficActionDrop
 		default:
-			util.WriteError(status.Errorf(status.InvalidArgument, "unknown action type"), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "unknown action type"), w)
 			return
 		}
 
-		switch r.Protocol {
+		switch rule.Protocol {
 		case api.PolicyRuleUpdateProtocolAll:
 			pr.Protocol = server.PolicyRuleProtocolALL
 		case api.PolicyRuleUpdateProtocolTcp:
@@ -171,14 +171,14 @@ func (h *Policies) savePolicy(
 		case api.PolicyRuleUpdateProtocolIcmp:
 			pr.Protocol = server.PolicyRuleProtocolICMP
 		default:
-			util.WriteError(status.Errorf(status.InvalidArgument, "unknown protocol type: %v", r.Protocol), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "unknown protocol type: %v", rule.Protocol), w)
 			return
 		}
 
-		if r.Ports != nil && len(*r.Ports) != 0 {
-			for _, v := range *r.Ports {
+		if rule.Ports != nil && len(*rule.Ports) != 0 {
+			for _, v := range *rule.Ports {
 				if port, err := strconv.Atoi(v); err != nil || port < 1 || port > 65535 {
-					util.WriteError(status.Errorf(status.InvalidArgument, "valid port value is in 1..65535 range"), w)
+					util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "valid port value is in 1..65535 range"), w)
 					return
 				}
 				pr.Ports = append(pr.Ports, v)
@@ -189,16 +189,16 @@ func (h *Policies) savePolicy(
 		switch pr.Protocol {
 		case server.PolicyRuleProtocolALL, server.PolicyRuleProtocolICMP:
 			if len(pr.Ports) != 0 {
-				util.WriteError(status.Errorf(status.InvalidArgument, "for ALL or ICMP protocol ports is not allowed"), w)
+				util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "for ALL or ICMP protocol ports is not allowed"), w)
 				return
 			}
 			if !pr.Bidirectional {
-				util.WriteError(status.Errorf(status.InvalidArgument, "for ALL or ICMP protocol type flow can be only bi-directional"), w)
+				util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "for ALL or ICMP protocol type flow can be only bi-directional"), w)
 				return
 			}
 		case server.PolicyRuleProtocolTCP, server.PolicyRuleProtocolUDP:
 			if !pr.Bidirectional && len(pr.Ports) == 0 {
-				util.WriteError(status.Errorf(status.InvalidArgument, "for ALL or ICMP protocol type flow can be only bi-directional"), w)
+				util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "for ALL or ICMP protocol type flow can be only bi-directional"), w)
 				return
 			}
 		}
@@ -210,26 +210,26 @@ func (h *Policies) savePolicy(
 		policy.SourcePostureChecks = sourcePostureChecksToStrings(account, *req.SourcePostureChecks)
 	}
 
-	if err := h.accountManager.SavePolicy(account.Id, user.Id, &policy); err != nil {
-		util.WriteError(err, w)
+	if err := h.accountManager.SavePolicy(r.Context(), account.Id, user.Id, &policy); err != nil {
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	resp := toPolicyResponse(account, &policy)
 	if len(resp.Rules) == 0 {
-		util.WriteError(status.Errorf(status.Internal, "no rules in the policy"), w)
+		util.WriteError(r.Context(), status.Errorf(status.Internal, "no rules in the policy"), w)
 		return
 	}
 
-	util.WriteJSONObject(w, resp)
+	util.WriteJSONObject(r.Context(), w, resp)
 }
 
 // DeletePolicy handles policy deletion request
 func (h *Policies) DeletePolicy(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 	aID := account.Id
@@ -237,24 +237,24 @@ func (h *Policies) DeletePolicy(w http.ResponseWriter, r *http.Request) {
 	vars := mux.Vars(r)
 	policyID := vars["policyId"]
 	if len(policyID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid policy ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid policy ID"), w)
 		return
 	}
 
-	if err = h.accountManager.DeletePolicy(aID, policyID, user.Id); err != nil {
-		util.WriteError(err, w)
+	if err = h.accountManager.DeletePolicy(r.Context(), aID, policyID, user.Id); err != nil {
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 // GetPolicy handles a group Get request identified by ID
 func (h *Policies) GetPolicy(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -263,25 +263,25 @@ func (h *Policies) GetPolicy(w http.ResponseWriter, r *http.Request) {
 		vars := mux.Vars(r)
 		policyID := vars["policyId"]
 		if len(policyID) == 0 {
-			util.WriteError(status.Errorf(status.InvalidArgument, "invalid policy ID"), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid policy ID"), w)
 			return
 		}
 
-		policy, err := h.accountManager.GetPolicy(account.Id, policyID, user.Id)
+		policy, err := h.accountManager.GetPolicy(r.Context(), account.Id, policyID, user.Id)
 		if err != nil {
-			util.WriteError(err, w)
+			util.WriteError(r.Context(), err, w)
 			return
 		}
 
 		resp := toPolicyResponse(account, policy)
 		if len(resp.Rules) == 0 {
-			util.WriteError(status.Errorf(status.Internal, "no rules in the policy"), w)
+			util.WriteError(r.Context(), status.Errorf(status.Internal, "no rules in the policy"), w)
 			return
 		}
 
-		util.WriteJSONObject(w, resp)
+		util.WriteJSONObject(r.Context(), w, resp)
 	default:
-		util.WriteError(status.Errorf(status.NotFound, "method not found"), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "method not found"), w)
 	}
 }
 

management/server/http/policies_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -30,21 +31,21 @@ func initPoliciesTestData(policies ...*server.Policy) *Policies {
 	}
 	return &Policies{
 		accountManager: &mock_server.MockAccountManager{
-			GetPolicyFunc: func(_, policyID, _ string) (*server.Policy, error) {
+			GetPolicyFunc: func(_ context.Context, _, policyID, _ string) (*server.Policy, error) {
 				policy, ok := testPolicies[policyID]
 				if !ok {
 					return nil, status.Errorf(status.NotFound, "policy not found")
 				}
 				return policy, nil
 			},
-			SavePolicyFunc: func(_, _ string, policy *server.Policy) error {
+			SavePolicyFunc: func(_ context.Context, _, _ string, policy *server.Policy) error {
 				if !strings.HasPrefix(policy.ID, "id-") {
 					policy.ID = "id-was-set"
 					policy.Rules[0].ID = "id-was-set"
 				}
 				return nil
 			},
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				user := server.NewAdminUser("test_user")
 				return &server.Account{
 					Id:     claims.AccountId,

management/server/http/posture_checks_handler.go

@@ -37,15 +37,15 @@ func NewPostureChecksHandler(accountManager server.AccountManager, geolocationMa
 // GetAllPostureChecks list for the account
 func (p *PostureChecksHandler) GetAllPostureChecks(w http.ResponseWriter, r *http.Request) {
 	claims := p.claimsExtractor.FromRequestContext(r)
-	account, user, err := p.accountManager.GetAccountFromToken(claims)
+	account, user, err := p.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	accountPostureChecks, err := p.accountManager.ListPostureChecks(account.Id, user.Id)
+	accountPostureChecks, err := p.accountManager.ListPostureChecks(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -54,22 +54,22 @@ func (p *PostureChecksHandler) GetAllPostureChecks(w http.ResponseWriter, r *htt
 		postureChecks = append(postureChecks, postureCheck.ToAPIResponse())
 	}
 
-	util.WriteJSONObject(w, postureChecks)
+	util.WriteJSONObject(r.Context(), w, postureChecks)
 }
 
 // UpdatePostureCheck handles update to a posture check identified by a given ID
 func (p *PostureChecksHandler) UpdatePostureCheck(w http.ResponseWriter, r *http.Request) {
 	claims := p.claimsExtractor.FromRequestContext(r)
-	account, user, err := p.accountManager.GetAccountFromToken(claims)
+	account, user, err := p.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	postureChecksID := vars["postureCheckId"]
 	if len(postureChecksID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid posture checks ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid posture checks ID"), w)
 		return
 	}
 
@@ -81,7 +81,7 @@ func (p *PostureChecksHandler) UpdatePostureCheck(w http.ResponseWriter, r *http
 		}
 	}
 	if postureChecksIdx < 0 {
-		util.WriteError(status.Errorf(status.NotFound, "couldn't find posture checks id %s", postureChecksID), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "couldn't find posture checks id %s", postureChecksID), w)
 		return
 	}
 
@@ -91,9 +91,9 @@ func (p *PostureChecksHandler) UpdatePostureCheck(w http.ResponseWriter, r *http
 // CreatePostureCheck handles posture check creation request
 func (p *PostureChecksHandler) CreatePostureCheck(w http.ResponseWriter, r *http.Request) {
 	claims := p.claimsExtractor.FromRequestContext(r)
-	account, user, err := p.accountManager.GetAccountFromToken(claims)
+	account, user, err := p.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -103,50 +103,50 @@ func (p *PostureChecksHandler) CreatePostureCheck(w http.ResponseWriter, r *http
 // GetPostureCheck handles a posture check Get request identified by ID
 func (p *PostureChecksHandler) GetPostureCheck(w http.ResponseWriter, r *http.Request) {
 	claims := p.claimsExtractor.FromRequestContext(r)
-	account, user, err := p.accountManager.GetAccountFromToken(claims)
+	account, user, err := p.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	postureChecksID := vars["postureCheckId"]
 	if len(postureChecksID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid posture checks ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid posture checks ID"), w)
 		return
 	}
 
-	postureChecks, err := p.accountManager.GetPostureChecks(account.Id, postureChecksID, user.Id)
+	postureChecks, err := p.accountManager.GetPostureChecks(r.Context(), account.Id, postureChecksID, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, postureChecks.ToAPIResponse())
+	util.WriteJSONObject(r.Context(), w, postureChecks.ToAPIResponse())
 }
 
 // DeletePostureCheck handles posture check deletion request
 func (p *PostureChecksHandler) DeletePostureCheck(w http.ResponseWriter, r *http.Request) {
 	claims := p.claimsExtractor.FromRequestContext(r)
-	account, user, err := p.accountManager.GetAccountFromToken(claims)
+	account, user, err := p.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	postureChecksID := vars["postureCheckId"]
 	if len(postureChecksID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid posture checks ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid posture checks ID"), w)
 		return
 	}
 
-	if err = p.accountManager.DeletePostureChecks(account.Id, postureChecksID, user.Id); err != nil {
-		util.WriteError(err, w)
+	if err = p.accountManager.DeletePostureChecks(r.Context(), account.Id, postureChecksID, user.Id); err != nil {
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 // savePostureChecks handles posture checks create and update
@@ -169,7 +169,7 @@ func (p *PostureChecksHandler) savePostureChecks(
 
 	if geoLocationCheck := req.Checks.GeoLocationCheck; geoLocationCheck != nil {
 		if p.geolocationManager == nil {
-			util.WriteError(status.Errorf(status.PreconditionFailed, "Geo location database is not initialized. "+
+			util.WriteError(r.Context(), status.Errorf(status.PreconditionFailed, "Geo location database is not initialized. "+
 				"Check the self-hosted Geo database documentation at https://docs.netbird.io/selfhosted/geo-support"), w)
 			return
 		}
@@ -177,14 +177,14 @@ func (p *PostureChecksHandler) savePostureChecks(
 
 	postureChecks, err := posture.NewChecksFromAPIPostureCheckUpdate(req, postureChecksID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	if err := p.accountManager.SavePostureChecks(account.Id, user.Id, postureChecks); err != nil {
-		util.WriteError(err, w)
+	if err := p.accountManager.SavePostureChecks(r.Context(), account.Id, user.Id, postureChecks); err != nil {
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, postureChecks.ToAPIResponse())
+	util.WriteJSONObject(r.Context(), w, postureChecks.ToAPIResponse())
 }

management/server/http/posture_checks_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -33,14 +34,14 @@ func initPostureChecksTestData(postureChecks ...*posture.Checks) *PostureChecksH
 
 	return &PostureChecksHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetPostureChecksFunc: func(accountID, postureChecksID, userID string) (*posture.Checks, error) {
+			GetPostureChecksFunc: func(_ context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error) {
 				p, ok := testPostureChecks[postureChecksID]
 				if !ok {
 					return nil, status.Errorf(status.NotFound, "posture checks not found")
 				}
 				return p, nil
 			},
-			SavePostureChecksFunc: func(accountID, userID string, postureChecks *posture.Checks) error {
+			SavePostureChecksFunc: func(_ context.Context, accountID, userID string, postureChecks *posture.Checks) error {
 				postureChecks.ID = "postureCheck"
 				testPostureChecks[postureChecks.ID] = postureChecks
 
@@ -50,7 +51,7 @@ func initPostureChecksTestData(postureChecks ...*posture.Checks) *PostureChecksH
 
 				return nil
 			},
-			DeletePostureChecksFunc: func(accountID, postureChecksID, userID string) error {
+			DeletePostureChecksFunc: func(_ context.Context, accountID, postureChecksID, userID string) error {
 				_, ok := testPostureChecks[postureChecksID]
 				if !ok {
 					return status.Errorf(status.NotFound, "posture checks not found")
@@ -59,14 +60,14 @@ func initPostureChecksTestData(postureChecks ...*posture.Checks) *PostureChecksH
 
 				return nil
 			},
-			ListPostureChecksFunc: func(accountID, userID string) ([]*posture.Checks, error) {
+			ListPostureChecksFunc: func(_ context.Context, accountID, userID string) ([]*posture.Checks, error) {
 				accountPostureChecks := make([]*posture.Checks, len(testPostureChecks))
 				for _, p := range testPostureChecks {
 					accountPostureChecks = append(accountPostureChecks, p)
 				}
 				return accountPostureChecks, nil
 			},
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				user := server.NewAdminUser("test_user")
 				return &server.Account{
 					Id: claims.AccountId,

management/server/http/routes_handler.go

@@ -43,36 +43,36 @@ func NewRoutesHandler(accountManager server.AccountManager, authCfg AuthCfg) *Ro
 // GetAllRoutes returns the list of routes for the account
 func (h *RoutesHandler) GetAllRoutes(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	routes, err := h.accountManager.ListRoutes(account.Id, user.Id)
+	routes, err := h.accountManager.ListRoutes(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 	apiRoutes := make([]*api.Route, 0)
-	for _, r := range routes {
-		route, err := toRouteResponse(r)
+	for _, route := range routes {
+		route, err := toRouteResponse(route)
 		if err != nil {
-			util.WriteError(status.Errorf(status.Internal, failedToConvertRoute, err), w)
+			util.WriteError(r.Context(), status.Errorf(status.Internal, failedToConvertRoute, err), w)
 			return
 		}
 		apiRoutes = append(apiRoutes, route)
 	}
 
-	util.WriteJSONObject(w, apiRoutes)
+	util.WriteJSONObject(r.Context(), w, apiRoutes)
 }
 
 // CreateRoute handles route creation request
 func (h *RoutesHandler) CreateRoute(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -84,7 +84,7 @@ func (h *RoutesHandler) CreateRoute(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if err := h.validateRoute(req); err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -94,7 +94,7 @@ func (h *RoutesHandler) CreateRoute(w http.ResponseWriter, r *http.Request) {
 	if req.Domains != nil {
 		d, err := validateDomains(*req.Domains)
 		if err != nil {
-			util.WriteError(status.Errorf(status.InvalidArgument, "invalid domains: %v", err), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid domains: %v", err), w)
 			return
 		}
 		domains = d
@@ -102,7 +102,7 @@ func (h *RoutesHandler) CreateRoute(w http.ResponseWriter, r *http.Request) {
 	} else if req.Network != nil {
 		networkType, newPrefix, err = route.ParseNetwork(*req.Network)
 		if err != nil {
-			util.WriteError(err, w)
+			util.WriteError(r.Context(), err, w)
 			return
 		}
 	}
@@ -120,24 +120,24 @@ func (h *RoutesHandler) CreateRoute(w http.ResponseWriter, r *http.Request) {
 	// Do not allow non-Linux peers
 	if peer := account.GetPeer(peerId); peer != nil {
 		if peer.Meta.GoOS != "linux" {
-			util.WriteError(status.Errorf(status.InvalidArgument, "non-linux peers are not supported as network routes"), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "non-linux peers are not supported as network routes"), w)
 			return
 		}
 	}
 
-	newRoute, err := h.accountManager.CreateRoute(account.Id, newPrefix, networkType, domains, peerId, peerGroupIds, req.Description, route.NetID(req.NetworkId), req.Masquerade, req.Metric, req.Groups, req.Enabled, user.Id, req.KeepRoute)
+	newRoute, err := h.accountManager.CreateRoute(r.Context(), account.Id, newPrefix, networkType, domains, peerId, peerGroupIds, req.Description, route.NetID(req.NetworkId), req.Masquerade, req.Metric, req.Groups, req.Enabled, user.Id, req.KeepRoute)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	routes, err := toRouteResponse(newRoute)
 	if err != nil {
-		util.WriteError(status.Errorf(status.Internal, failedToConvertRoute, err), w)
+		util.WriteError(r.Context(), status.Errorf(status.Internal, failedToConvertRoute, err), w)
 		return
 	}
 
-	util.WriteJSONObject(w, routes)
+	util.WriteJSONObject(r.Context(), w, routes)
 }
 
 func (h *RoutesHandler) validateRoute(req api.PostApiRoutesJSONRequestBody) error {
@@ -168,22 +168,22 @@ func (h *RoutesHandler) validateRoute(req api.PostApiRoutesJSONRequestBody) erro
 // UpdateRoute handles update to a route identified by a given ID
 func (h *RoutesHandler) UpdateRoute(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	routeID := vars["routeId"]
 	if len(routeID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid route ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid route ID"), w)
 		return
 	}
 
-	_, err = h.accountManager.GetRoute(account.Id, route.ID(routeID), user.Id)
+	_, err = h.accountManager.GetRoute(r.Context(), account.Id, route.ID(routeID), user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -195,7 +195,7 @@ func (h *RoutesHandler) UpdateRoute(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if err := h.validateRoute(req); err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -207,7 +207,7 @@ func (h *RoutesHandler) UpdateRoute(w http.ResponseWriter, r *http.Request) {
 	// do not allow non Linux peers
 	if peer := account.GetPeer(peerID); peer != nil {
 		if peer.Meta.GoOS != "linux" {
-			util.WriteError(status.Errorf(status.InvalidArgument, "non-linux peers are non supported as network routes"), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "non-linux peers are non supported as network routes"), w)
 			return
 		}
 	}
@@ -226,7 +226,7 @@ func (h *RoutesHandler) UpdateRoute(w http.ResponseWriter, r *http.Request) {
 	if req.Domains != nil {
 		d, err := validateDomains(*req.Domains)
 		if err != nil {
-			util.WriteError(status.Errorf(status.InvalidArgument, "invalid domains: %v", err), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid domains: %v", err), w)
 			return
 		}
 		newRoute.Domains = d
@@ -234,7 +234,7 @@ func (h *RoutesHandler) UpdateRoute(w http.ResponseWriter, r *http.Request) {
 	} else if req.Network != nil {
 		newRoute.NetworkType, newRoute.Network, err = route.ParseNetwork(*req.Network)
 		if err != nil {
-			util.WriteError(err, w)
+			util.WriteError(r.Context(), err, w)
 			return
 		}
 	}
@@ -247,73 +247,73 @@ func (h *RoutesHandler) UpdateRoute(w http.ResponseWriter, r *http.Request) {
 		newRoute.PeerGroups = *req.PeerGroups
 	}
 
-	err = h.accountManager.SaveRoute(account.Id, user.Id, newRoute)
+	err = h.accountManager.SaveRoute(r.Context(), account.Id, user.Id, newRoute)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	routes, err := toRouteResponse(newRoute)
 	if err != nil {
-		util.WriteError(status.Errorf(status.Internal, failedToConvertRoute, err), w)
+		util.WriteError(r.Context(), status.Errorf(status.Internal, failedToConvertRoute, err), w)
 		return
 	}
 
-	util.WriteJSONObject(w, routes)
+	util.WriteJSONObject(r.Context(), w, routes)
 }
 
 // DeleteRoute handles route deletion request
 func (h *RoutesHandler) DeleteRoute(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	routeID := mux.Vars(r)["routeId"]
 	if len(routeID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid route ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid route ID"), w)
 		return
 	}
 
-	err = h.accountManager.DeleteRoute(account.Id, route.ID(routeID), user.Id)
+	err = h.accountManager.DeleteRoute(r.Context(), account.Id, route.ID(routeID), user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 // GetRoute handles a route Get request identified by ID
 func (h *RoutesHandler) GetRoute(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	routeID := mux.Vars(r)["routeId"]
 	if len(routeID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid route ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid route ID"), w)
 		return
 	}
 
-	foundRoute, err := h.accountManager.GetRoute(account.Id, route.ID(routeID), user.Id)
+	foundRoute, err := h.accountManager.GetRoute(r.Context(), account.Id, route.ID(routeID), user.Id)
 	if err != nil {
-		util.WriteError(status.Errorf(status.NotFound, "route not found"), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "route not found"), w)
 		return
 	}
 
 	routes, err := toRouteResponse(foundRoute)
 	if err != nil {
-		util.WriteError(status.Errorf(status.Internal, failedToConvertRoute, err), w)
+		util.WriteError(r.Context(), status.Errorf(status.Internal, failedToConvertRoute, err), w)
 		return
 	}
 
-	util.WriteJSONObject(w, routes)
+	util.WriteJSONObject(r.Context(), w, routes)
 }
 
 func toRouteResponse(serverRoute *route.Route) (*api.Route, error) {

management/server/http/routes_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -89,7 +90,7 @@ var testingAccount = &server.Account{
 func initRoutesTestData() *RoutesHandler {
 	return &RoutesHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetRouteFunc: func(_ string, routeID route.ID, _ string) (*route.Route, error) {
+			GetRouteFunc: func(_ context.Context, _ string, routeID route.ID, _ string) (*route.Route, error) {
 				if routeID == existingRouteID {
 					return baseExistingRoute, nil
 				}
@@ -104,7 +105,7 @@ func initRoutesTestData() *RoutesHandler {
 				}
 				return nil, status.Errorf(status.NotFound, "route with ID %s not found", routeID)
 			},
-			CreateRouteFunc: func(accountID string, prefix netip.Prefix, networkType route.NetworkType, domains domain.List, peerID string, peerGroups []string, description string, netID route.NetID, masquerade bool, metric int, groups []string, enabled bool, _ string, keepRoute bool) (*route.Route, error) {
+			CreateRouteFunc: func(_ context.Context, accountID string, prefix netip.Prefix, networkType route.NetworkType, domains domain.List, peerID string, peerGroups []string, description string, netID route.NetID, masquerade bool, metric int, groups []string, enabled bool, _ string, keepRoute bool) (*route.Route, error) {
 				if peerID == notFoundPeerID {
 					return nil, status.Errorf(status.InvalidArgument, "peer with ID %s not found", peerID)
 				}
@@ -126,19 +127,19 @@ func initRoutesTestData() *RoutesHandler {
 					KeepRoute:   keepRoute,
 				}, nil
 			},
-			SaveRouteFunc: func(_, _ string, r *route.Route) error {
+			SaveRouteFunc: func(_ context.Context, _, _ string, r *route.Route) error {
 				if r.Peer == notFoundPeerID {
 					return status.Errorf(status.InvalidArgument, "peer with ID %s not found", r.Peer)
 				}
 				return nil
 			},
-			DeleteRouteFunc: func(_ string, routeID route.ID, _ string) error {
+			DeleteRouteFunc: func(_ context.Context, _ string, routeID route.ID, _ string) error {
 				if routeID != existingRouteID {
 					return status.Errorf(status.NotFound, "Peer with ID %s not found", routeID)
 				}
 				return nil
 			},
-			GetAccountFromTokenFunc: func(_ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, _ jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return testingAccount, testingAccount.Users["test_user"], nil
 			},
 		},

management/server/http/setupkeys_handler.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"time"
@@ -34,9 +35,9 @@ func NewSetupKeysHandler(accountManager server.AccountManager, authCfg AuthCfg)
 // CreateSetupKey is a POST requests that creates a new SetupKey
 func (h *SetupKeysHandler) CreateSetupKey(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -48,13 +49,13 @@ func (h *SetupKeysHandler) CreateSetupKey(w http.ResponseWriter, r *http.Request
 	}
 
 	if req.Name == "" {
-		util.WriteError(status.Errorf(status.InvalidArgument, "setup key name shouldn't be empty"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "setup key name shouldn't be empty"), w)
 		return
 	}
 
 	if !(server.SetupKeyType(req.Type) == server.SetupKeyReusable ||
 		server.SetupKeyType(req.Type) == server.SetupKeyOneOff) {
-		util.WriteError(status.Errorf(status.InvalidArgument, "unknown setup key type %s", req.Type), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "unknown setup key type %s", req.Type), w)
 		return
 	}
 
@@ -63,7 +64,7 @@ func (h *SetupKeysHandler) CreateSetupKey(w http.ResponseWriter, r *http.Request
 	day := time.Hour * 24
 	year := day * 365
 	if expiresIn < day || expiresIn > year {
-		util.WriteError(status.Errorf(status.InvalidArgument, "expiresIn should be between 1 day and 365 days"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "expiresIn should be between 1 day and 365 days"), w)
 		return
 	}
 
@@ -75,54 +76,54 @@ func (h *SetupKeysHandler) CreateSetupKey(w http.ResponseWriter, r *http.Request
 	if req.Ephemeral != nil {
 		ephemeral = *req.Ephemeral
 	}
-	setupKey, err := h.accountManager.CreateSetupKey(account.Id, req.Name, server.SetupKeyType(req.Type), expiresIn,
+	setupKey, err := h.accountManager.CreateSetupKey(r.Context(), account.Id, req.Name, server.SetupKeyType(req.Type), expiresIn,
 		req.AutoGroups, req.UsageLimit, user.Id, ephemeral)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	writeSuccess(w, setupKey)
+	writeSuccess(r.Context(), w, setupKey)
 }
 
 // GetSetupKey is a GET request to get a SetupKey by ID
 func (h *SetupKeysHandler) GetSetupKey(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	keyID := vars["keyId"]
 	if len(keyID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid key ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid key ID"), w)
 		return
 	}
 
-	key, err := h.accountManager.GetSetupKey(account.Id, user.Id, keyID)
+	key, err := h.accountManager.GetSetupKey(r.Context(), account.Id, user.Id, keyID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	writeSuccess(w, key)
+	writeSuccess(r.Context(), w, key)
 }
 
 // UpdateSetupKey is a PUT request to update server.SetupKey
 func (h *SetupKeysHandler) UpdateSetupKey(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	keyID := vars["keyId"]
 	if len(keyID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid key ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid key ID"), w)
 		return
 	}
 
@@ -134,12 +135,12 @@ func (h *SetupKeysHandler) UpdateSetupKey(w http.ResponseWriter, r *http.Request
 	}
 
 	if req.Name == "" {
-		util.WriteError(status.Errorf(status.InvalidArgument, "setup key name field is invalid: %s", req.Name), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "setup key name field is invalid: %s", req.Name), w)
 		return
 	}
 
 	if req.AutoGroups == nil {
-		util.WriteError(status.Errorf(status.InvalidArgument, "setup key AutoGroups field is invalid"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "setup key AutoGroups field is invalid"), w)
 		return
 	}
 
@@ -149,26 +150,26 @@ func (h *SetupKeysHandler) UpdateSetupKey(w http.ResponseWriter, r *http.Request
 	newKey.Name = req.Name
 	newKey.Id = keyID
 
-	newKey, err = h.accountManager.SaveSetupKey(account.Id, newKey, user.Id)
+	newKey, err = h.accountManager.SaveSetupKey(r.Context(), account.Id, newKey, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
-	writeSuccess(w, newKey)
+	writeSuccess(r.Context(), w, newKey)
 }
 
 // GetAllSetupKeys is a GET request that returns a list of SetupKey
 func (h *SetupKeysHandler) GetAllSetupKeys(w http.ResponseWriter, r *http.Request) {
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	setupKeys, err := h.accountManager.ListSetupKeys(account.Id, user.Id)
+	setupKeys, err := h.accountManager.ListSetupKeys(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -177,15 +178,15 @@ func (h *SetupKeysHandler) GetAllSetupKeys(w http.ResponseWriter, r *http.Reques
 		apiSetupKeys = append(apiSetupKeys, toResponseBody(key))
 	}
 
-	util.WriteJSONObject(w, apiSetupKeys)
+	util.WriteJSONObject(r.Context(), w, apiSetupKeys)
 }
 
-func writeSuccess(w http.ResponseWriter, key *server.SetupKey) {
+func writeSuccess(ctx context.Context, w http.ResponseWriter, key *server.SetupKey) {
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(200)
 	err := json.NewEncoder(w).Encode(toResponseBody(key))
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(ctx, err, w)
 		return
 	}
 }

management/server/http/setupkeys_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -33,7 +34,7 @@ func initSetupKeysTestMetaData(defaultKey *server.SetupKey, newKey *server.Setup
 ) *SetupKeysHandler {
 	return &SetupKeysHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return &server.Account{
 					Id:     testAccountID,
 					Domain: "hotmail.com",
@@ -49,7 +50,7 @@ func initSetupKeysTestMetaData(defaultKey *server.SetupKey, newKey *server.Setup
 					},
 				}, user, nil
 			},
-			CreateSetupKeyFunc: func(_ string, keyName string, typ server.SetupKeyType, _ time.Duration, _ []string,
+			CreateSetupKeyFunc: func(_ context.Context, _ string, keyName string, typ server.SetupKeyType, _ time.Duration, _ []string,
 				_ int, _ string, ephemeral bool,
 			) (*server.SetupKey, error) {
 				if keyName == newKey.Name || typ != newKey.Type {
@@ -59,7 +60,7 @@ func initSetupKeysTestMetaData(defaultKey *server.SetupKey, newKey *server.Setup
 				}
 				return nil, fmt.Errorf("failed creating setup key")
 			},
-			GetSetupKeyFunc: func(accountID, userID, keyID string) (*server.SetupKey, error) {
+			GetSetupKeyFunc: func(_ context.Context, accountID, userID, keyID string) (*server.SetupKey, error) {
 				switch keyID {
 				case defaultKey.Id:
 					return defaultKey, nil
@@ -70,14 +71,14 @@ func initSetupKeysTestMetaData(defaultKey *server.SetupKey, newKey *server.Setup
 				}
 			},
 
-			SaveSetupKeyFunc: func(accountID string, key *server.SetupKey, _ string) (*server.SetupKey, error) {
+			SaveSetupKeyFunc: func(_ context.Context, accountID string, key *server.SetupKey, _ string) (*server.SetupKey, error) {
 				if key.Id == updatedSetupKey.Id {
 					return updatedSetupKey, nil
 				}
 				return nil, status.Errorf(status.NotFound, "key %s not found", key.Id)
 			},
 
-			ListSetupKeysFunc: func(accountID, userID string) ([]*server.SetupKey, error) {
+			ListSetupKeysFunc: func(_ context.Context, accountID, userID string) ([]*server.SetupKey, error) {
 				return []*server.SetupKey{defaultKey}, nil
 			},
 		},

management/server/http/users_handler.go

@@ -41,22 +41,22 @@ func (h *UsersHandler) UpdateUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	userID := vars["userId"]
 	if len(userID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
 	existingUser, ok := account.Users[userID]
 	if !ok {
-		util.WriteError(status.Errorf(status.NotFound, "couldn't find user with ID %s", userID), w)
+		util.WriteError(r.Context(), status.Errorf(status.NotFound, "couldn't find user with ID %s", userID), w)
 		return
 	}
 
@@ -74,11 +74,11 @@ func (h *UsersHandler) UpdateUser(w http.ResponseWriter, r *http.Request) {
 
 	userRole := server.StrRoleToUserRole(req.Role)
 	if userRole == server.UserRoleUnknown {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user role"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user role"), w)
 		return
 	}
 
-	newUser, err := h.accountManager.SaveUser(account.Id, user.Id, &server.User{
+	newUser, err := h.accountManager.SaveUser(r.Context(), account.Id, user.Id, &server.User{
 		Id:                   userID,
 		Role:                 userRole,
 		AutoGroups:           req.AutoGroups,
@@ -88,10 +88,10 @@ func (h *UsersHandler) UpdateUser(w http.ResponseWriter, r *http.Request) {
 	})
 
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
-	util.WriteJSONObject(w, toUserResponse(newUser, claims.UserId))
+	util.WriteJSONObject(r.Context(), w, toUserResponse(newUser, claims.UserId))
 }
 
 // DeleteUser is a DELETE request to delete a user
@@ -102,26 +102,26 @@ func (h *UsersHandler) DeleteUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	targetUserID := vars["userId"]
 	if len(targetUserID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
-	err = h.accountManager.DeleteUser(account.Id, user.Id, targetUserID)
+	err = h.accountManager.DeleteUser(r.Context(), account.Id, user.Id, targetUserID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 // CreateUser creates a User in the system with a status "invited" (effectively this is a user invite).
@@ -132,9 +132,9 @@ func (h *UsersHandler) CreateUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
@@ -146,7 +146,7 @@ func (h *UsersHandler) CreateUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if server.StrRoleToUserRole(req.Role) == server.UserRoleUnknown {
-		util.WriteError(status.Errorf(status.InvalidArgument, "unknown user role %s", req.Role), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "unknown user role %s", req.Role), w)
 		return
 	}
 
@@ -160,7 +160,7 @@ func (h *UsersHandler) CreateUser(w http.ResponseWriter, r *http.Request) {
 		name = *req.Name
 	}
 
-	newUser, err := h.accountManager.CreateUser(account.Id, user.Id, &server.UserInfo{
+	newUser, err := h.accountManager.CreateUser(r.Context(), account.Id, user.Id, &server.UserInfo{
 		Email:         email,
 		Name:          name,
 		Role:          req.Role,
@@ -169,10 +169,10 @@ func (h *UsersHandler) CreateUser(w http.ResponseWriter, r *http.Request) {
 		Issued:        server.UserIssuedAPI,
 	})
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
-	util.WriteJSONObject(w, toUserResponse(newUser, claims.UserId))
+	util.WriteJSONObject(r.Context(), w, toUserResponse(newUser, claims.UserId))
 }
 
 // GetAllUsers returns a list of users of the account this user belongs to.
@@ -184,42 +184,42 @@ func (h *UsersHandler) GetAllUsers(w http.ResponseWriter, r *http.Request) {
 	}
 
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	data, err := h.accountManager.GetUsersFromAccount(account.Id, user.Id)
+	data, err := h.accountManager.GetUsersFromAccount(r.Context(), account.Id, user.Id)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	serviceUser := r.URL.Query().Get("service_user")
 
 	users := make([]*api.User, 0)
-	for _, r := range data {
-		if r.NonDeletable {
+	for _, d := range data {
+		if d.NonDeletable {
 			continue
 		}
 		if serviceUser == "" {
-			users = append(users, toUserResponse(r, claims.UserId))
+			users = append(users, toUserResponse(d, claims.UserId))
 			continue
 		}
 
 		includeServiceUser, err := strconv.ParseBool(serviceUser)
-		log.Debugf("Should include service user: %v", includeServiceUser)
+		log.WithContext(r.Context()).Debugf("Should include service user: %v", includeServiceUser)
 		if err != nil {
-			util.WriteError(status.Errorf(status.InvalidArgument, "invalid service_user query parameter"), w)
+			util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid service_user query parameter"), w)
 			return
 		}
-		if includeServiceUser == r.IsServiceUser {
-			users = append(users, toUserResponse(r, claims.UserId))
+		if includeServiceUser == d.IsServiceUser {
+			users = append(users, toUserResponse(d, claims.UserId))
 		}
 	}
 
-	util.WriteJSONObject(w, users)
+	util.WriteJSONObject(r.Context(), w, users)
 }
 
 // InviteUser resend invitations to users who haven't activated their accounts,
@@ -231,26 +231,26 @@ func (h *UsersHandler) InviteUser(w http.ResponseWriter, r *http.Request) {
 	}
 
 	claims := h.claimsExtractor.FromRequestContext(r)
-	account, user, err := h.accountManager.GetAccountFromToken(claims)
+	account, user, err := h.accountManager.GetAccountFromToken(r.Context(), claims)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
 	vars := mux.Vars(r)
 	targetUserID := vars["userId"]
 	if len(targetUserID) == 0 {
-		util.WriteError(status.Errorf(status.InvalidArgument, "invalid user ID"), w)
+		util.WriteError(r.Context(), status.Errorf(status.InvalidArgument, "invalid user ID"), w)
 		return
 	}
 
-	err = h.accountManager.InviteUser(account.Id, user.Id, targetUserID)
+	err = h.accountManager.InviteUser(r.Context(), account.Id, user.Id, targetUserID)
 	if err != nil {
-		util.WriteError(err, w)
+		util.WriteError(r.Context(), err, w)
 		return
 	}
 
-	util.WriteJSONObject(w, emptyObject{})
+	util.WriteJSONObject(r.Context(), w, emptyObject{})
 }
 
 func toUserResponse(user *server.UserInfo, currenUserID string) *api.User {

management/server/http/users_handler_test.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -63,10 +64,10 @@ var usersTestAccount = &server.Account{
 func initUsersTestData() *UsersHandler {
 	return &UsersHandler{
 		accountManager: &mock_server.MockAccountManager{
-			GetAccountFromTokenFunc: func(claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
+			GetAccountFromTokenFunc: func(_ context.Context, claims jwtclaims.AuthorizationClaims) (*server.Account, *server.User, error) {
 				return usersTestAccount, usersTestAccount.Users[claims.UserId], nil
 			},
-			GetUsersFromAccountFunc: func(accountID, userID string) ([]*server.UserInfo, error) {
+			GetUsersFromAccountFunc: func(_ context.Context, accountID, userID string) ([]*server.UserInfo, error) {
 				users := make([]*server.UserInfo, 0)
 				for _, v := range usersTestAccount.Users {
 					users = append(users, &server.UserInfo{
@@ -81,13 +82,13 @@ func initUsersTestData() *UsersHandler {
 				}
 				return users, nil
 			},
-			CreateUserFunc: func(accountID, userID string, key *server.UserInfo) (*server.UserInfo, error) {
+			CreateUserFunc: func(_ context.Context, accountID, userID string, key *server.UserInfo) (*server.UserInfo, error) {
 				if userID != existingUserID {
 					return nil, status.Errorf(status.NotFound, "user with ID %s does not exists", userID)
 				}
 				return key, nil
 			},
-			DeleteUserFunc: func(accountID string, initiatorUserID string, targetUserID string) error {
+			DeleteUserFunc: func(_ context.Context, accountID string, initiatorUserID string, targetUserID string) error {
 				if targetUserID == notFoundUserID {
 					return status.Errorf(status.NotFound, "user with ID %s does not exists", targetUserID)
 				}
@@ -96,7 +97,7 @@ func initUsersTestData() *UsersHandler {
 				}
 				return nil
 			},
-			SaveUserFunc: func(accountID, userID string, update *server.User) (*server.UserInfo, error) {
+			SaveUserFunc: func(_ context.Context, accountID, userID string, update *server.User) (*server.UserInfo, error) {
 				if update.Id == notFoundUserID {
 					return nil, status.Errorf(status.NotFound, "user with ID %s does not exists", update.Id)
 				}
@@ -111,7 +112,7 @@ func initUsersTestData() *UsersHandler {
 				}
 				return info, nil
 			},
-			InviteUserFunc: func(accountID string, initiatorUserID string, targetUserID string) error {
+			InviteUserFunc: func(_ context.Context, accountID string, initiatorUserID string, targetUserID string) error {
 				if initiatorUserID != existingUserID {
 					return status.Errorf(status.NotFound, "user with ID %s does not exists", initiatorUserID)
 				}

management/server/http/util/util.go

@@ -1,6 +1,7 @@
 package util
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -19,12 +20,12 @@ type ErrorResponse struct {
 }
 
 // WriteJSONObject simply writes object to the HTTP response in JSON format
-func WriteJSONObject(w http.ResponseWriter, obj interface{}) {
+func WriteJSONObject(ctx context.Context, w http.ResponseWriter, obj interface{}) {
 	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
 	w.WriteHeader(http.StatusOK)
 	err := json.NewEncoder(w).Encode(obj)
 	if err != nil {
-		WriteError(err, w)
+		WriteError(ctx, err, w)
 		return
 	}
 }
@@ -76,8 +77,8 @@ func WriteErrorResponse(errMsg string, httpStatus int, w http.ResponseWriter) {
 
 // WriteError converts an error to an JSON error response.
 // If it is known internal error of type server.Error then it sets the messages from the error, a generic message otherwise
-func WriteError(err error, w http.ResponseWriter) {
-	log.Errorf("got a handler error: %s", err.Error())
+func WriteError(ctx context.Context, err error, w http.ResponseWriter) {
+	log.WithContext(ctx).Errorf("got a handler error: %s", err.Error())
 	errStatus, ok := status.FromError(err)
 	httpStatus := http.StatusInternalServerError
 	msg := "internal server error"
@@ -106,7 +107,7 @@ func WriteError(err error, w http.ResponseWriter) {
 		msg = strings.ToLower(err.Error())
 	} else {
 		unhandledMSG := fmt.Sprintf("got unhandled error code, error: %s", err.Error())
-		log.Error(unhandledMSG)
+		log.WithContext(ctx).Error(unhandledMSG)
 	}
 
 	WriteErrorResponse(msg, httpStatus, w)