Fix(auth0) caching Users by accountId

* Update links in Start using NetBird

* Update internals overview and co structure

* Netbird to NetBird

client/Dockerfile

@@ -1,4 +1,4 @@
 FROM gcr.io/distroless/base:debug
 ENV WT_LOG_FILE=console
 ENTRYPOINT [ "/go/bin/netbird","up"]
-COPY netbird /go/bin/netbird
+COPY netbird /go/bin/netbird

management/server/account.go

@@ -328,11 +328,12 @@ func (am *DefaultAccountManager) GetUsersFromAccount(accountID string) ([]*UserI
 
 	queriedUsers := make([]*idp.UserData, 0)
 	if !isNil(am.idpManager) {
-		queriedUsers, err = am.idpManager.GetBatchedUserData(accountID)
+		queriedUsers, err = am.idpManager.GetAllUsers(accountID)
 		if err != nil {
 			return nil, err
 		}
 	}
+	// TODO: we need to check whether we need to refresh our cache or not
 
 	userInfo := make([]*UserInfo, 0)
 
@@ -352,7 +353,6 @@ func (am *DefaultAccountManager) GetUsersFromAccount(accountID string) ([]*UserI
 	for _, queriedUser := range queriedUsers {
 		if localUser, contains := account.Users[queriedUser.ID]; contains {
 			userInfo = append(userInfo, mergeLocalAndQueryUser(*queriedUser, *localUser))
-			log.Debugf("Merged userinfo to send back; %v", userInfo)
 		}
 	}
 

management/server/grpcserver.go

@@ -203,8 +203,11 @@ func (s *Server) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest) (*Pe
 		},
 	})
 	if err != nil {
-		if s, ok := status.FromError(err); ok && s.Code() == codes.FailedPrecondition {
-			return nil, err
+		s, ok := status.FromError(err)
+		if ok {
+			if s.Code() == codes.FailedPrecondition || s.Code() == codes.OutOfRange {
+				return nil, err
+			}
 		}
 		return nil, status.Errorf(codes.NotFound, "provided setup key doesn't exists")
 	}

management/server/idp/auth0.go

@@ -1,6 +1,9 @@
 package idp
 
 import (
+	"bytes"
+	"compress/gzip"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -18,10 +21,11 @@ import (
 
 // Auth0Manager auth0 manager client instance
 type Auth0Manager struct {
-	authIssuer  string
-	httpClient  ManagerHTTPClient
-	credentials ManagerCredentials
-	helper      ManagerHelper
+	authIssuer             string
+	httpClient             ManagerHTTPClient
+	credentials            ManagerCredentials
+	helper                 ManagerHelper
+	cachedUsersByAccountId map[string][]Auth0Profile
 }
 
 // Auth0ClientConfig auth0 manager client configurations
@@ -51,6 +55,38 @@ type Auth0Credentials struct {
 	mux          sync.Mutex
 }
 
+type Auth0Profile struct {
+	AccountId string `json:"wt_account_id"`
+	UserID    string `json:"user_id"`
+	Name      string `json:"name"`
+	Email     string `json:"email"`
+	CreatedAt string `json:"created_at"`
+	LastLogin string `json:"last_login"`
+}
+
+type UserExportJobResponse struct {
+	Type         string    `json:"type"`
+	Status       string    `json:"status"`
+	ConnectionId string    `json:"connection_id"`
+	Format       string    `json:"format"`
+	Limit        int       `json:"limit"`
+	Connection   string    `json:"connection"`
+	CreatedAt    time.Time `json:"created_at"`
+	Id           string    `json:"id"`
+}
+
+type ExportJobStatusResponse struct {
+	Type         string    `json:"type"`
+	Status       string    `json:"status"`
+	ConnectionId string    `json:"connection_id"`
+	Format       string    `json:"format"`
+	Limit        int       `json:"limit"`
+	Location     string    `json:"location"`
+	Connection   string    `json:"connection"`
+	CreatedAt    time.Time `json:"created_at"`
+	Id           string    `json:"id"`
+}
+
 // NewAuth0Manager creates a new instance of the Auth0Manager
 func NewAuth0Manager(config Auth0ClientConfig) (*Auth0Manager, error) {
 
@@ -81,11 +117,13 @@ func NewAuth0Manager(config Auth0ClientConfig) (*Auth0Manager, error) {
 		httpClient:   httpClient,
 		helper:       helper,
 	}
+
 	return &Auth0Manager{
-		authIssuer:  config.AuthIssuer,
-		credentials: credentials,
-		httpClient:  httpClient,
-		helper:      helper,
+		authIssuer:             config.AuthIssuer,
+		credentials:            credentials,
+		httpClient:             httpClient,
+		helper:                 helper,
+		cachedUsersByAccountId: make(map[string][]Auth0Profile),
 	}, nil
 }
 
@@ -186,44 +224,198 @@ func (c *Auth0Credentials) Authenticate() (JWTToken, error) {
 	return c.jwtToken, nil
 }
 
-func batchRequestUsersUrl(authIssuer, accountId string, page int) (string, url.Values, error) {
-	u, err := url.Parse(authIssuer + "/api/v2/users")
-	if err != nil {
-		return "", nil, err
-	}
-	q := u.Query()
-	q.Set("page", strconv.Itoa(page))
-	q.Set("search_engine", "v3")
-	q.Set("q", "app_metadata.wt_account_id:"+accountId)
-	u.RawQuery = q.Encode()
-
-	return u.String(), q, nil
-}
-
-func requestByUserIdUrl(authIssuer, userId string) string {
-	return authIssuer + "/api/v2/users/" + userId
-}
-
-// GetBatchedUserData requests users in batches from Auth0
-func (am *Auth0Manager) GetBatchedUserData(accountId string) ([]*UserData, error) {
-	jwtToken, err := am.credentials.Authenticate()
-	if err != nil {
-		return nil, err
+// Gets all users from cache, if the cache exists
+// Otherwise we will initialize the cache with creating the export job on auth0
+func (am *Auth0Manager) GetAllUsers(accountId string) ([]*UserData, error) {
+	if len(am.cachedUsersByAccountId[accountId]) == 0 {
+		err := am.createExportUsersJob(accountId)
+		if err != nil {
+			log.Debugf("Couldn't cache users; %v", err)
+			return nil, err
+		}
 	}
 
 	var list []*UserData
 
+	cachedUsers := am.cachedUsersByAccountId[accountId]
+	for _, val := range cachedUsers {
+		list = append(list, &UserData{
+			Name:  val.Name,
+			Email: val.Email,
+			ID:    val.UserID,
+		})
+	}
+
+	return list, nil
+}
+
+// This creates an export job on auth0 for all users.
+func (am *Auth0Manager) createExportUsersJob(accountId string) error {
+	jwtToken, err := am.credentials.Authenticate()
+	if err != nil {
+		return err
+	}
+
+	reqURL := am.authIssuer + "/api/v2/jobs/users-exports"
+
+	payloadString := fmt.Sprintf("{\"format\": \"json\"," +
+		"\"fields\": [{\"name\": \"created_at\"}, {\"name\": \"last_login\"},{\"name\": \"user_id\"}, {\"name\": \"email\"}, {\"name\": \"name\"}, {\"name\": \"app_metadata.wt_account_id\", \"export_as\": \"wt_account_id\"}]}")
+
+	payload := strings.NewReader(payloadString)
+
+	exportJobReq, err := http.NewRequest("POST", reqURL, payload)
+	if err != nil {
+		return err
+	}
+	exportJobReq.Header.Add("authorization", "Bearer "+jwtToken.AccessToken)
+	exportJobReq.Header.Add("content-type", "application/json")
+
+	jobResp, err := am.httpClient.Do(exportJobReq)
+	if err != nil {
+		log.Debugf("Couldn't get job response %v", err)
+		return err
+	}
+
+	defer func() {
+		err = jobResp.Body.Close()
+		if err != nil {
+			log.Errorf("error while closing update user app metadata response body: %v", err)
+		}
+	}()
+	if jobResp.StatusCode != 200 {
+		return fmt.Errorf("unable to update the appMetadata, statusCode %d", jobResp.StatusCode)
+	}
+
+	var exportJobResp UserExportJobResponse
+
+	body, err := ioutil.ReadAll(jobResp.Body)
+	if err != nil {
+		log.Debugf("Coudln't read export job response; %v", err)
+		return err
+	}
+
+	err = am.helper.Unmarshal(body, &exportJobResp)
+	if err != nil {
+		log.Debugf("Coudln't unmarshal export job response; %v", err)
+		return err
+	}
+
+	if exportJobResp.Id == "" {
+		return fmt.Errorf("couldn't get an batch id status %d, %s, response body: %v", jobResp.StatusCode, jobResp.Status, exportJobResp)
+	}
+
+	log.Debugf("batch id status %d, %s, response body: %v", jobResp.StatusCode, jobResp.Status, exportJobResp)
+
+	ctx, cancel := context.WithTimeout(context.TODO(), 90*time.Second)
+	defer cancel()
+
+	done, downloadLink, err := am.checkExportJobStatus(ctx, exportJobResp.Id)
+	if err != nil {
+		log.Debugf("Failed at getting status checks from exportJob; %v", err)
+		return err
+	}
+
+	if done {
+		err = am.cacheUsers(downloadLink)
+		if err != nil {
+			log.Debugf("Failed to cache users via download link; %v", err)
+		}
+	}
+
+	return nil
+}
+
+// Downloads the users from auth0 and caches it in memory
+// Users are only cached if they have an wt_account_id stored in auth0
+func (am *Auth0Manager) cacheUsers(location string) error {
+	body, err := doGetReq(am.httpClient, location, "")
+	if err != nil {
+		log.Debugf("Can't download cached users; %v", err)
+		return err
+	}
+
+	bodyReader := bytes.NewReader(body)
+
+	gzipReader, err := gzip.NewReader(bodyReader)
+	if err != nil {
+		return err
+	}
+
+	decoder := json.NewDecoder(gzipReader)
+
+	for decoder.More() {
+		profile := Auth0Profile{}
+		err = decoder.Decode(&profile)
+		if err != nil {
+			log.Errorf("Couldn't decode profile; %v", err)
+			return err
+		}
+		if profile.AccountId != "" {
+			am.cachedUsersByAccountId[profile.AccountId] = append(am.cachedUsersByAccountId[profile.AccountId], profile)
+		}
+	}
+
+	return nil
+}
+
+// This checks the status of the job created at CreateExportUsersJob.
+// If the status is "completed", then return the downloadLink
+func (am *Auth0Manager) checkExportJobStatus(ctx context.Context, jobId string) (bool, string, error) {
+	retry := time.NewTicker(time.Second)
+	for {
+		select {
+		case <-ctx.Done():
+			log.Debugf("Export job status stopped...\n")
+			return false, "", ctx.Err()
+		case <-retry.C:
+			jwtToken, err := am.credentials.Authenticate()
+			if err != nil {
+				return false, "", err
+			}
+
+			statusUrl := am.authIssuer + "/api/v2/jobs/" + jobId
+			body, err := doGetReq(am.httpClient, statusUrl, jwtToken.AccessToken)
+			if err != nil {
+				return false, "", err
+			}
+
+			var status ExportJobStatusResponse
+			err = am.helper.Unmarshal(body, &status)
+			if err != nil {
+				return false, "", err
+			}
+
+			log.Debugf("Current export job status is %v", status.Status)
+
+			if status.Status != "completed" {
+				continue
+			}
+
+			return true, status.Location, nil
+		}
+	}
+}
+
+// Invalidates old cache for Account and re-queries it from auth0
+func (am *Auth0Manager) forceUpdateUserCache(accountId string) error {
+	jwtToken, err := am.credentials.Authenticate()
+	if err != nil {
+		return err
+	}
+
+	var list []Auth0Profile
+
 	// https://auth0.com/docs/manage-users/user-search/retrieve-users-with-get-users-endpoint#limitations
 	// auth0 limitation of 1000 users via this endpoint
 	for page := 0; page < 20; page++ {
 		reqURL, query, err := batchRequestUsersUrl(am.authIssuer, accountId, page)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		req, err := http.NewRequest(http.MethodGet, reqURL, strings.NewReader(query.Encode()))
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		req.Header.Add("authorization", "Bearer "+jwtToken.AccessToken)
@@ -231,41 +423,42 @@ func (am *Auth0Manager) GetBatchedUserData(accountId string) ([]*UserData, error
 
 		res, err := am.httpClient.Do(req)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		body, err := io.ReadAll(res.Body)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
-		var batch []UserData
+		var batch []Auth0Profile
 		err = json.Unmarshal(body, &batch)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		log.Debugf("requested batch; %v", batch)
 
 		err = res.Body.Close()
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		if res.StatusCode != 200 {
-			return nil, fmt.Errorf("unable to request UserData from auth0, statusCode %d", res.StatusCode)
+			return fmt.Errorf("unable to request UserData from auth0, statusCode %d", res.StatusCode)
 		}
 
 		if len(batch) == 0 {
-			return list, nil
+			return nil
 		}
 
 		for user := range batch {
-			list = append(list, &batch[user])
+			list = append(list, batch[user])
 		}
 	}
+	am.cachedUsersByAccountId[accountId] = list
 
-	return list, nil
+	return nil
 }
 
 // GetUserDataByID requests user data from auth0 via ID
@@ -359,3 +552,54 @@ func (am *Auth0Manager) UpdateUserAppMetadata(userId string, appMetadata AppMeta
 
 	return nil
 }
+
+func batchRequestUsersUrl(authIssuer, accountId string, page int) (string, url.Values, error) {
+	u, err := url.Parse(authIssuer + "/api/v2/users")
+	if err != nil {
+		return "", nil, err
+	}
+	q := u.Query()
+	q.Set("page", strconv.Itoa(page))
+	q.Set("search_engine", "v3")
+	q.Set("q", "app_metadata.wt_account_id:"+accountId)
+	u.RawQuery = q.Encode()
+
+	return u.String(), q, nil
+}
+
+func requestByUserIdUrl(authIssuer, userId string) string {
+	return authIssuer + "/api/v2/users/" + userId
+}
+
+// Boilerplate implementation for Get Requests.
+func doGetReq(client ManagerHTTPClient, url, accessToken string) ([]byte, error) {
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	if accessToken != "" {
+		req.Header.Add("authorization", "Bearer "+accessToken)
+	}
+
+	res, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	defer func() {
+		err = res.Body.Close()
+		if err != nil {
+			log.Errorf("error while closing body for url %s: %v", url, err)
+		}
+	}()
+	if res.StatusCode != 200 {
+		return nil, fmt.Errorf("unable to get %s, statusCode %d", url, res.StatusCode)
+	}
+
+	body, err := ioutil.ReadAll(res.Body)
+	if err != nil {
+		return nil, err
+	}
+	return body, nil
+}

management/server/idp/idp.go

@@ -11,7 +11,7 @@ import (
 type Manager interface {
 	UpdateUserAppMetadata(userId string, appMetadata AppMetadata) error
 	GetUserDataByID(userId string, appMetadata AppMetadata) (*UserData, error)
-	GetBatchedUserData(accountId string) ([]*UserData, error)
+	GetAllUsers(accountId string) ([]*UserData, error)
 }
 
 // Config an idp configuration struct to be loaded from management server's config file

management/server/network.go

@@ -1,9 +1,10 @@
 package server
 
 import (
-	"fmt"
 	"github.com/c-robinson/iplib"
 	"github.com/rs/xid"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"math/rand"
 	"net"
 	"sync"
@@ -77,13 +78,10 @@ func AllocatePeerIP(ipNet net.IPNet, takenIps []net.IP) (net.IP, error) {
 		takenIPMap[ip.String()] = struct{}{}
 	}
 
-	ips, _, err := generateIPs(&ipNet, takenIPMap)
-	if err != nil {
-		return nil, fmt.Errorf("failed allocating new IP for the ipNet %s and takenIps %s", ipNet.String(), takenIps)
-	}
+	ips, _ := generateIPs(&ipNet, takenIPMap)
 
 	if len(ips) == 0 {
-		return nil, fmt.Errorf("failed allocating new IP for the ipNet %s - network is out of IPs", ipNet.String())
+		return nil, status.Errorf(codes.OutOfRange, "failed allocating new IP for the ipNet %s - network is out of IPs", ipNet.String())
 	}
 
 	// pick a random IP
@@ -95,7 +93,7 @@ func AllocatePeerIP(ipNet net.IPNet, takenIps []net.IP) (net.IP, error) {
 }
 
 // generateIPs generates a list of all possible IPs of the given network excluding IPs specified in the exclusion list
-func generateIPs(ipNet *net.IPNet, exclusions map[string]struct{}) ([]net.IP, int, error) {
+func generateIPs(ipNet *net.IPNet, exclusions map[string]struct{}) ([]net.IP, int) {
 
 	var ips []net.IP
 	for ip := ipNet.IP.Mask(ipNet.Mask); ipNet.Contains(ip); incIP(ip) {
@@ -108,10 +106,10 @@ func generateIPs(ipNet *net.IPNet, exclusions map[string]struct{}) ([]net.IP, in
 	lenIPs := len(ips)
 	switch {
 	case lenIPs < 2:
-		return ips, lenIPs, nil
+		return ips, lenIPs
 
 	default:
-		return ips[1 : len(ips)-1], lenIPs - 2, nil
+		return ips[1 : len(ips)-1], lenIPs - 2
 	}
 }
 

management/server/peer.go

@@ -346,7 +346,10 @@ func (am *DefaultAccountManager) AddPeer(
 	}
 
 	network := account.Network
-	nextIp, _ := AllocatePeerIP(network.Net, takenIps)
+	nextIp, err := AllocatePeerIP(network.Net, takenIps)
+	if err != nil {
+		return nil, err
+	}
 
 	newPeer := &Peer{
 		Key:      peer.Key,