teleport/lib/msgraph/client.go

// Teleport
// Copyright (C) 2025 Gravitational, Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package msgraph

import (
	"bytes"
	"cmp"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/google/uuid"
	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"

	"github.com/gravitational/teleport"
	apidefaults "github.com/gravitational/teleport/api/defaults"
	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/api/utils/retryutils"
	"github.com/gravitational/teleport/lib/defaults"
	"github.com/gravitational/teleport/lib/msgraph/models"
	"github.com/gravitational/teleport/lib/observability/metrics"
	"github.com/gravitational/teleport/lib/utils"
)

// graphVersion is the default version of the MS Graph API endpoint.
const graphVersion = "v1.0"

// defaultPageSize is the page size used when [Config.PageSize] is not specified.
const defaultPageSize = 500

// scopes defines OAuth scopes the client authenticates for.
var scopes = []string{"https://graph.microsoft.com/.default"}

// AzureTokenProvider defines a method to get an authorization token from the Entra STS.
// Concrete implementations of this are defined by [github.com/Azure/azure-sdk-for-go/sdk/azidentity].
type AzureTokenProvider interface {
	// GetToken requests an access token from Microsoft Entra ID. Token providers from azidentity
	// return cached tokens whenever possible and are safe for concurrent use.
	// https://github.com/Azure/azure-sdk-for-go/blob/sdk/azidentity/v1.11.0/sdk/azidentity/TOKEN_CACHING.MD
	GetToken(ctx context.Context, opts policy.TokenRequestOptions) (azcore.AccessToken, error)
}

// DeltaStore defines an interface for persisting delta links.
type DeltaStore interface {
	// Get returns a delta link for the given endpoint.
	Get(endpoint string) string
	// Set sets a delta link for the given endpoint.
	Set(endpoint, deltaLink string)
	// Clear deletes delta link for the given endpoint.
	Clear(endpoint string)
}

func defaultHTTPClient() (*http.Client, error) {
	transport, err := defaults.Transport()
	if err != nil {
		return nil, trace.Wrap(err)
	}

	transport.ExpectContinueTimeout = apidefaults.DefaultIOTimeout
	transport.ResponseHeaderTimeout = apidefaults.DefaultIOTimeout
	transport.IdleConnTimeout = apidefaults.DefaultIdleTimeout

	return &http.Client{
		Transport: transport,
		Timeout:   apidefaults.DefaultIOTimeout,
	}, nil
}

// Config defines configuration options for [client].
type Config struct {
	// TokenProvider provides tokens to authorize to MS Graph API.
	TokenProvider AzureTokenProvider
	// HTTPClient is the HTTP client to use for calls to the API.
	// If not specified, [http.DefaultClient] is used.
	HTTPClient *http.Client
	// Clock is the clock to use for time operations (e.g. delay when retrying requests).
	Clock clockwork.Clock
	// RetryConfig specifies parameters for retrying failed requests.
	// Client will prefer to use the `Retry-After` header returned from the API,
	// and only use this retry config if the header is not provided.
	RetryConfig *retryutils.RetryV2Config
	// PageSize limits the number of objects to return in one batch when using paginated requests (via the `$top` parameter).
	PageSize int
	// GraphEndpoint specifies root domain of the Graph API.
	GraphEndpoint string
	Logger        *slog.Logger
	// MetricsRegistry configures where metrics should be registered.
	// When nil, metrics are created but not registered.
	MetricsRegistry *metrics.Registry
}

// SetDefaults sets the default values for optional fields.
func (cfg *Config) SetDefaults() {
	defaultHTTPClient, _ := defaultHTTPClient()

	cfg.HTTPClient = cmp.Or(cfg.HTTPClient, defaultHTTPClient)
	cfg.Clock = cmp.Or(cfg.Clock, clockwork.NewRealClock())
	cfg.RetryConfig = cmp.Or(cfg.RetryConfig, &retryutils.RetryV2Config{
		First:  1 * time.Second,
		Driver: retryutils.NewExponentialDriver(1 * time.Second),
		Max:    defaults.HighResPollingPeriod,
	})
	if cfg.PageSize <= 0 {
		cfg.PageSize = defaultPageSize
	}
	if cfg.GraphEndpoint == "" {
		cfg.GraphEndpoint = types.MSGraphDefaultEndpoint
	}
	if cfg.Logger == nil {
		cfg.Logger = slog.With(teleport.ComponentKey, "msgraph")
	}
	if cfg.MetricsRegistry == nil {
		cfg.MetricsRegistry = metrics.NoopRegistry()
	}
}

// Validate checks that required fields are set.
func (cfg *Config) Validate() error {
	if cfg.TokenProvider == nil {
		return trace.BadParameter("TokenProvider must be set")
	}
	if cfg.HTTPClient == nil {
		return trace.BadParameter("HTTPClient must be set")
	}
	if err := types.ValidateMSGraphEndpoint(cfg.GraphEndpoint); err != nil {
		return trace.Wrap(err)
	}
	return nil
}

type Client struct {
	httpClient    *http.Client
	tokenProvider AzureTokenProvider
	clock         clockwork.Clock
	retryConfig   retryutils.RetryV2Config
	baseURL       *url.URL
	pageSize      int
	logger        *slog.Logger
	metrics       *clientMetrics
}

// NewClient returns a new client for the given config.
func NewClient(cfg Config) (*Client, error) {
	cfg.SetDefaults()
	if err := cfg.Validate(); err != nil {
		return nil, trace.Wrap(err)
	}
	base, err := url.Parse(cfg.GraphEndpoint)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	m := newMetrics(cfg.MetricsRegistry)
	// gracefully handle not being given a metric registry
	if err := m.register(cfg.MetricsRegistry); err != nil {
		cfg.Logger.ErrorContext(context.Background(), "Failed to register metrics.", "error", err)
	}

	return &Client{
		httpClient:    cfg.HTTPClient,
		tokenProvider: cfg.TokenProvider,
		clock:         cfg.Clock,
		retryConfig:   *cfg.RetryConfig,
		baseURL:       base.JoinPath(graphVersion),
		pageSize:      cfg.PageSize,
		logger:        cfg.Logger,
		metrics:       m,
	}, nil
}

// request is the base function for HTTP API calls.
// It implements retry handling in case of API throttling, see [https://learn.microsoft.com/en-us/graph/throttling].
// If the response from the Graph API has status code outside of [200, 400) range, request attempts
// to parse the response body as [GraphError] and if successful returns it as error.
func (c *Client) request(ctx context.Context, method string, uri string, header http.Header, payload []byte) (*http.Response, error) {
	var body io.ReadSeeker = nil
	if len(payload) > 0 {
		body = bytes.NewReader(payload)
	}

	req, err := http.NewRequestWithContext(ctx, method, uri, body)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	for key := range header {
		for _, value := range header.Values(key) {
			req.Header.Add(key, value)
		}
	}

	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}

	const maxRetries = 5
	var retryAfter time.Duration

	// RetryV2 only used when the API does not return a Retry-After header.
	retry, err := retryutils.NewRetryV2(c.retryConfig)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	var lastErr error
	var start time.Time
	for i := range maxRetries {
		if retryAfter > 0 {
			select {
			case <-c.clock.After(retryAfter):
			case <-ctx.Done():
				return nil, trace.NewAggregate(ctx.Err(), trace.Wrap(lastErr, "%s %s", req.Method, req.URL.Path))
			}
		}
		token, err := c.tokenProvider.GetToken(ctx, policy.TokenRequestOptions{
			Scopes: scopes,
		})
		if err != nil {
			authFailedError := &azidentity.AuthenticationFailedError{}
			if ok := errors.As(err, &authFailedError); ok && authFailedError.RawResponse != nil &&
				authFailedError.RawResponse.Body != nil {
				resp := authFailedError.RawResponse
				authError, conversionErr := readAuthError(resp.Body, resp.StatusCode)
				resp.Body.Close()
				if conversionErr == nil {
					err = authError
				}
			}
			return nil, trace.Wrap(err, "failed to get azure authentication token")
		}
		req.Header.Set("Authorization", "Bearer "+token.Token)

		requestID := uuid.NewString()
		// https://learn.microsoft.com/en-us/graph/best-practices-concept#reliability-and-support
		req.Header.Set("client-request-id", requestID)

		start = c.clock.Now()
		resp, err := c.httpClient.Do(req)
		if err != nil {
			return nil, trace.Wrap(err) // hard I/O error, bail
		}
		c.metrics.requestDuration.WithLabelValues(method).Observe(c.clock.Since(start).Seconds())
		c.metrics.requestTotal.WithLabelValues(method, strconv.Itoa(resp.StatusCode))

		if resp.StatusCode >= 200 && resp.StatusCode < 400 {
			return resp, nil
		}

		respBody, err := utils.ReadAtMost(resp.Body, teleport.MaxHTTPResponseSize)
		if err != nil {
			return nil, trace.Wrap(err)
		}
		if err := resp.Body.Close(); err != nil {
			c.logger.WarnContext(req.Context(), "Failed to close http.Response body", "error", err)
		}

		c.logger.DebugContext(req.Context(), "Request failed",
			"body", string(respBody),
			"status", resp.StatusCode,
			"url", req.URL,
			"client_request_id", requestID,
			"retry_count", i,
		)

		retryAfter = retry.Duration()
		retryAfterFromHeader := time.Duration(0)
		if ra := resp.Header.Get("Retry-After"); ra != "" {
			if seconds, err := strconv.Atoi(ra); err != nil {
				c.logger.WarnContext(ctx, `Failed to parse "Retry-After" header`, "error", err)
			} else {
				retryAfterFromHeader = time.Duration(seconds) * time.Second
				retryAfter = time.Duration(seconds) * time.Second
			}
		}
		retry.Inc()

		graphError, err := readError(respBody, resp.StatusCode)
		if err != nil {
			lastErr = err // error while reading the graph error, relay
		} else if graphError != nil {
			if retryAfterFromHeader > 0 {
				graphError.RetryAfter = retryAfterFromHeader
			}
			lastErr = trace.Wrap(graphError)
		} else {
			// API did not return a valid error structure, best-effort reporting.
			lastErr = trace.Errorf("%s", resp.Status)
		}
		if !isRetriable(resp.StatusCode) {
			break
		}

		// prepare for the next request attempt by rewinding the body
		if body != nil {
			_, err := body.Seek(0, io.SeekStart)
			if err != nil {
				return nil, trace.Wrap(err)
			}
		}
	}
	return nil, trace.Wrap(lastErr, "%s %s", req.Method, req.URL.Path)
}

func (c *Client) endpointURI(segments ...string) *url.URL {
	escapedSegments := make([]string, 0, cap(segments))
	for _, s := range segments {
		// Handling of slash vs escaped slash (%2F) in paths is ambiguous and inconsistent.
		// See e.g.: https://stackoverflow.com/questions/1957115/is-a-slash-equivalent-to-an-encoded-slash-2f-in-the-path-portion-of-a
		// We do not expect slashes to be needed within a single path segment,
		// so we just remove slashes from each segment.
		escapedSegments = append(escapedSegments, url.PathEscape(strings.ReplaceAll(s, "/", "")))
	}
	uri := c.baseURL
	uri = uri.JoinPath(escapedSegments...)
	return uri
}

// roundtrip makes a request to the API,
// serializing `in` as a JSON body, and deserializing the response as the given type `T`.
// It is used for GET and POST requests, where a response body is expected.
func roundtrip[T any](ctx context.Context, c *Client, method string, uri string, in any) (T, error) {
	var zero T
	var body []byte
	var err error
	if in != nil {
		body, err = json.Marshal(in)
		if err != nil {
			return zero, trace.Wrap(err)
		}
	}
	resp, err := c.request(ctx, method, uri, nil /* extra headers */, body)
	if err != nil {
		return zero, trace.Wrap(err)
	}
	defer resp.Body.Close()

	var out T
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return zero, trace.Wrap(err)
	}
	return out, nil
}

// patch makes a PATCH request to the API, serializing `in` as a JSON body.
// It expects a 204 No Content response.
func (c *Client) patch(ctx context.Context, uri string, in any) error {
	body, err := json.Marshal(in)
	if err != nil {
		return trace.Wrap(err)
	}
	resp, err := c.request(ctx, http.MethodPatch, uri, nil /* extra headers */, body)
	if err != nil {
		return trace.Wrap(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusNoContent {
		return trace.BadParameter("expected a 204 No Content response, got status code %v", resp.StatusCode)
	}
	return nil
}

// CreateFederatedIdentityCredential creates a new FederatedCredential.
// Ref: [https://learn.microsoft.com/en-us/graph/api/application-post-federatedidentitycredentials].
func (c *Client) CreateFederatedIdentityCredential(ctx context.Context, appObjectID string, cred *models.FederatedIdentityCredential) (*models.FederatedIdentityCredential, error) {
	uri := c.endpointURI("applications", appObjectID, "federatedIdentityCredentials")
	out, err := roundtrip[*models.FederatedIdentityCredential](ctx, c, http.MethodPost, uri.String(), cred)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// CreateServicePrincipalTokenSigningCertificate generates a new token signing certificate for the given service principal.
// Ref: [https://learn.microsoft.com/en-us/graph/api/serviceprincipal-addtokensigningcertificate].
func (c *Client) CreateServicePrincipalTokenSigningCertificate(ctx context.Context, spID string, displayName string) (*models.SelfSignedCertificate, error) {
	uri := c.endpointURI("servicePrincipals", spID, "addTokenSigningCertificate")
	in := map[string]string{"displayName": displayName}
	out, err := roundtrip[*models.SelfSignedCertificate](ctx, c, http.MethodPost, uri.String(), in)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// GetServicePrincipalByAppId returns the service principal associated with the given application.
// Note that appID here is the app the application "client ID" ([Application.AppID]), not "object ID" ([Application.ID]).
// Ref: [https://learn.microsoft.com/en-us/graph/api/serviceprincipal-get].
func (c *Client) GetServicePrincipalByAppId(ctx context.Context, appID string) (*models.ServicePrincipal, error) {
	uri := c.endpointURI(fmt.Sprintf("servicePrincipals(appId='%s')", appID))
	out, err := roundtrip[*models.ServicePrincipal](ctx, c, http.MethodGet, uri.String(), nil)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// GetServicePrincipalsByDisplayName returns the service principals that have the given display name.
// Ref: [https://learn.microsoft.com/en-us/graph/api/serviceprincipal-list].
func (c *Client) GetServicePrincipalsByDisplayName(ctx context.Context, displayName string) ([]*models.ServicePrincipal, error) {
	filter := fmt.Sprintf("displayName eq '%s'", displayName)
	uri := c.endpointURI("servicePrincipals")
	uri.RawQuery = url.Values{
		"$filter": {filter},
	}.Encode()
	out, err := roundtrip[oDataListResponse[*models.ServicePrincipal]](ctx, c, http.MethodGet, uri.String(), nil)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out.Value, nil
}

// GetServicePrincipal returns the service principal for the given principal ID.
// Ref: [https://learn.microsoft.com/en-us/graph/api/serviceprincipal-get].
func (c *Client) GetServicePrincipal(ctx context.Context, principalId string) (*models.ServicePrincipal, error) {
	uri := c.endpointURI(fmt.Sprintf("servicePrincipals/%s", principalId))
	out, err := roundtrip[*models.ServicePrincipal](ctx, c, http.MethodGet, uri.String(), nil)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// GrantAppRoleToServicePrincipal grants the given app role to the specified Service Principal.
// Ref: [https://learn.microsoft.com/en-us/graph/api/serviceprincipal-post-approleassignedto]
func (c *Client) GrantAppRoleToServicePrincipal(ctx context.Context, spID string, assignment *models.AppRoleAssignment) (*models.AppRoleAssignment, error) {
	uri := c.endpointURI("servicePrincipals", spID, "appRoleAssignedTo")
	out, err := roundtrip[*models.AppRoleAssignment](ctx, c, http.MethodPost, uri.String(), assignment)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// InstantiateApplicationTemplate instantiates an application from the Entra application Gallery,
// creating a pair of [Application] and [ServicePrincipal].
// Ref: [https://learn.microsoft.com/en-us/graph/api/applicationtemplate-instantiate].
func (c *Client) InstantiateApplicationTemplate(ctx context.Context, appTemplateID string, displayName string) (*models.ApplicationServicePrincipal, error) {
	uri := c.endpointURI("applicationTemplates", appTemplateID, "instantiate")
	in := map[string]string{
		"displayName": displayName,
	}
	out, err := roundtrip[*models.ApplicationServicePrincipal](ctx, c, http.MethodPost, uri.String(), in)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// UpdateApplication issues a partial update for an [Application].
// Note that appID here is the app the application  "object ID" ([Application.ID]), not "client ID" ([Application.AppID]).
// Ref: [https://learn.microsoft.com/en-us/graph/api/application-update].
func (c *Client) UpdateApplication(ctx context.Context, appObjectID string, app *models.Application) error {
	uri := c.endpointURI("applications", appObjectID)
	return trace.Wrap(c.patch(ctx, uri.String(), app))
}

// GetApplication returns the application with the given app client ID.
// Note that appID here is the app the application "client ID" ([Application.AppID]) not  "object ID" ([Application.ID]).
// Ref: [https://learn.microsoft.com/en-us/graph/api/application-get].
func (c *Client) GetApplication(ctx context.Context, applicationID string) (*models.Application, error) {
	applicationIDFilter := fmt.Sprintf("applications(appId='%s')", applicationID)
	uri := c.endpointURI(applicationIDFilter)
	out, err := roundtrip[*models.Application](ctx, c, http.MethodGet, uri.String(), nil)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	return out, nil
}

// UpdateServicePrincipal issues a partial update for a [ServicePrincipal].
// Ref: [https://learn.microsoft.com/en-us/graph/api/serviceprincipal-update].
func (c *Client) UpdateServicePrincipal(ctx context.Context, spID string, sp *models.ServicePrincipal) error {
	uri := c.endpointURI("servicePrincipals", spID)
	return trace.Wrap(c.patch(ctx, uri.String(), sp))
}

// isRetriable returns `true` when the given HTTP status code should be retried.
func isRetriable(code int) bool {
	return code == http.StatusTooManyRequests || code == http.StatusServiceUnavailable || code == http.StatusGatewayTimeout
}