[management] REST Client Account Impersonation Support

client/internal/dns/local/local.go

@@ -110,12 +110,8 @@ func (d *Resolver) Update(update []nbdns.SimpleRecord) {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
-	log.Infof("updating %d records. Records: %v", len(update), update)
-
 	maps.Clear(d.records)
 
-	log.Infof("map size: %d", len(d.records))
-
 	for _, rec := range update {
 		if err := d.registerRecord(rec); err != nil {
 			log.Warnf("failed to register the record (%s): %v", rec, err)

management/client/rest/accounts_test.go

@@ -66,6 +66,15 @@ func TestAccounts_List_Err(t *testing.T) {
 	})
 }
 
+func TestAccounts_List_ConnErr(t *testing.T) {
+	withMockClient(func(c *rest.Client, mux *http.ServeMux) {
+		ret, err := c.Accounts.List(context.Background())
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "404")
+		assert.Empty(t, ret)
+	})
+}
+
 func TestAccounts_Update_200(t *testing.T) {
 	withMockClient(func(c *rest.Client, mux *http.ServeMux) {
 		mux.HandleFunc("/api/accounts/Test", func(w http.ResponseWriter, r *http.Request) {

management/client/rest/client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 
@@ -12,9 +13,10 @@ import (
 
 // Client Management service HTTP REST API Client
 type Client struct {
-	managementURL string
-	authHeader    string
-	httpClient    HttpClient
+	managementURL       string
+	authHeader          string
+	impersonatedAccount string
+	httpClient          HTTPClient
 
 	// Accounts NetBird account APIs
 	// see more: https://docs.netbird.io/api/resources/accounts
@@ -85,7 +87,8 @@ func NewWithBearerToken(managementURL, token string) *Client {
 	)
 }
 
-func NewWithOptions(opts ...option) *Client {
+// NewWithOptions initialize new Client instance with options
+func NewWithOptions(opts ...Option) *Client {
 	client := &Client{
 		httpClient: http.DefaultClient,
 	}
@@ -114,6 +117,7 @@ func (c *Client) initialize() {
 	c.Events = &EventsAPI{c}
 }
 
+// NewRequest creates and executes new management API request
 func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
 	req, err := http.NewRequestWithContext(ctx, method, c.managementURL+path, body)
 	if err != nil {
@@ -126,6 +130,12 @@ func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Re
 		req.Header.Add("Content-Type", "application/json")
 	}
 
+	if c.impersonatedAccount != "" {
+		query := req.URL.Query()
+		query.Add("account", c.impersonatedAccount)
+		req.URL.RawQuery = query.Encode()
+	}
+
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		return nil, err
@@ -134,7 +144,8 @@ func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Re
 	if resp.StatusCode > 299 {
 		parsedErr, pErr := parseResponse[util.ErrorResponse](resp)
 		if pErr != nil {
-			return nil, err
+
+			return nil, pErr
 		}
 		return nil, errors.New(parsedErr.Message)
 	}
@@ -145,13 +156,16 @@ func (c *Client) NewRequest(ctx context.Context, method, path string, body io.Re
 func parseResponse[T any](resp *http.Response) (T, error) {
 	var ret T
 	if resp.Body == nil {
-		return ret, errors.New("No body")
+		return ret, fmt.Errorf("Body missing, HTTP Error code %d", resp.StatusCode)
 	}
 	bs, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return ret, err
 	}
 	err = json.Unmarshal(bs, &ret)
+	if err != nil {
+		return ret, fmt.Errorf("Error code %d, error unmarshalling body: %w", resp.StatusCode, err)
+	}
 
-	return ret, err
+	return ret, nil
 }

management/client/rest/options.go

@@ -2,34 +2,50 @@ package rest
 
 import "net/http"
 
-type option func(*Client)
+// Option modifier for creation of Client
+type Option func(*Client)
 
-type HttpClient interface {
+// HTTPClient interface for HTTP client
+type HTTPClient interface {
 	Do(req *http.Request) (*http.Response, error)
 }
 
-func WithHttpClient(client HttpClient) option {
+// WithHTTPClient overrides HTTPClient used
+func WithHTTPClient(client HTTPClient) Option {
 	return func(c *Client) {
 		c.httpClient = client
 	}
 }
 
-func WithBearerToken(token string) option {
+// WithBearerToken uses provided bearer token acquired from SSO for authentication
+func WithBearerToken(token string) Option {
 	return WithAuthHeader("Bearer " + token)
 }
 
-func WithPAT(token string) option {
+// WithPAT uses provided Personal Access Token
+// (created from NetBird Management Dashboard) for authentication
+func WithPAT(token string) Option {
 	return WithAuthHeader("Token " + token)
 }
 
-func WithManagementURL(url string) option {
+// WithManagementURL overrides target NetBird Management server
+func WithManagementURL(url string) Option {
 	return func(c *Client) {
 		c.managementURL = url
 	}
 }
 
-func WithAuthHeader(value string) option {
+// WithAuthHeader overrides auth header completely, this should generally not be used
+// and WithBearerToken or WithPAT should be used instead
+func WithAuthHeader(value string) Option {
 	return func(c *Client) {
 		c.authHeader = value
 	}
 }
+
+// WithAccount uses impersonated account for Client
+func WithAccount(value string) Option {
+	return func(c *Client) {
+		c.impersonatedAccount = value
+	}
+}