Permify / permify

internal/authn/oidc/authn.go

package oidc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/golang-jwt/jwt/v4"
	grpcauth "github.com/grpc-ecosystem/go-grpc-middleware/auth"
	"github.com/hashicorp/go-retryablehttp"
	"github.com/lestrrat-go/jwx/jwk"

	"github.com/Permify/permify/internal/config"
	base "github.com/Permify/permify/pkg/pb/base/v1"
)

type Authn struct {
	// URL of the issuer. This is typically the base URL of the identity provider.
	IssuerURL string
	// Audience for which the token is intended. It must match the audience in the JWT.
	Audience string
	// URL of the JSON Web Key Set (JWKS). This URL hosts public keys used to verify JWT signatures.
	JwksURI string
	// Pointer to an AutoRefresh object from the JWKS library. It helps in automatically refreshing the JWKS at predefined intervals.
	jwksSet *jwk.AutoRefresh
	// List of valid signing methods. Specifies which signing algorithms are considered valid for the JWTs.
	validMethods []string
	// Pointer to a JWT parser object. This is used to parse and validate the JWT tokens.
	jwtParser *jwt.Parser
	// Duration of the interval between retries for the backoff policy.
	backoffInterval time.Duration
	// Maximum number of retries for the backoff policy.
	backoffMaxRetries int

	backoffFrequency time.Duration

	// Global backoff state
	globalRetryCount  int
	globalFirstSeen   time.Time
	globalRetryKeyIds map[string]bool
	mu                sync.Mutex
}

// NewOidcAuthn initializes a new instance of the Authn struct with OpenID Connect (OIDC) configuration.
// It takes in a context for managing cancellation and a configuration object. It returns a pointer to an Authn instance or an error.
func NewOidcAuthn(ctx context.Context, conf config.Oidc) (*Authn, error) {
	// Create a new HTTP client with retry capabilities. This client is used for making HTTP requests, particularly for fetching OIDC configuration.
	client := retryablehttp.NewClient()
	client.Logger = SlogAdapter{Logger: slog.Default()}

	// Fetch the OIDC configuration from the issuer's well-known configuration endpoint.
	oidcConf, err := fetchOIDCConfiguration(client.StandardClient(), strings.TrimSuffix(conf.Issuer, "/")+"/.well-known/openid-configuration")
	if err != nil {
		// If there is an error fetching the OIDC configuration, return nil and the error.
		return nil, fmt.Errorf("failed to fetch OIDC configuration: %w", err)

unrecognized printf verb 'w'

	}

	// Set up automatic refresh of the JSON Web Key Set (JWKS) to ensure the public keys are always up-to-date.
	ar := jwk.NewAutoRefresh(ctx)                                                                                              // Create a new AutoRefresh instance for the JWKS.
	ar.Configure(oidcConf.JWKsURI, jwk.WithHTTPClient(client.StandardClient()), jwk.WithRefreshInterval(conf.RefreshInterval)) // Configure the auto-refresh parameters.

	// Validate and set backoffInterval, backoffMaxRetries, and backoffFrequency
	backoffInterval := conf.BackoffInterval
	if backoffInterval <= 0 {
		return nil, errors.New("invalid or missing backoffInterval")
	}

	backoffMaxRetries := conf.BackoffMaxRetries
	if backoffMaxRetries <= 0 {
		return nil, errors.New("invalid or missing backoffMaxRetries")
	}

	backoffFrequency := conf.BackoffFrequency
	if backoffFrequency <= 0 {
		return nil, errors.New("invalid or missing backoffFrequency")
	}

	// Initialize the Authn struct with the OIDC configuration details and other relevant settings.
	oidc := &Authn{
		IssuerURL:         conf.Issuer,                                            // URL of the token issuer.
		Audience:          conf.Audience,                                          // Intended audience of the token.
		JwksURI:           oidcConf.JWKsURI,                                       // URL of the JWKS endpoint.
		validMethods:      conf.ValidMethods,                                      // List of acceptable signing methods for the tokens.
		jwtParser:         jwt.NewParser(jwt.WithValidMethods(conf.ValidMethods)), // JWT parser configured with the valid signing methods.
		jwksSet:           ar,                                                     // Set the JWKS auto-refresh instance.
		backoffInterval:   backoffInterval,
		backoffMaxRetries: backoffMaxRetries,
		backoffFrequency:  backoffFrequency,
		globalRetryCount:  0,
		globalRetryKeyIds: make(map[string]bool),
		globalFirstSeen:   time.Time{},
		mu:                sync.Mutex{},
	}

	// Attempt to fetch the JWKS immediately to ensure it's available and valid.
	_, err = oidc.jwksSet.Fetch(ctx, oidc.JwksURI)
	if err != nil {
		// If there is an error fetching the JWKS, return nil and the error.
		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)

unrecognized printf verb 'w'

	}

	// Return the initialized OIDC authentication object and no error.
	return oidc, nil
}

// Authenticate validates the JWT token found in the authorization header of the incoming request.
// It uses the OIDC configuration to validate the token against the issuer's public keys.
func (oidc *Authn) Authenticate(requestContext context.Context) error {
	// Extract the authorization header from the metadata of the incoming gRPC request.
	authHeader, err := grpcauth.AuthFromMD(requestContext, "Bearer")
	if err != nil {
		// Log the error if the authorization header is missing or does not start with "Bearer"
		slog.Error("failed to extract authorization header from gRPC request", "error", err)
		// Return an error indicating the missing or incorrect bearer token
		return errors.New(base.ErrorCode_ERROR_CODE_MISSING_BEARER_TOKEN.String())
	}

	// Log the successful extraction of the authorization header for debugging purposes.
	// Remember, do not log the actual content of authHeader as it might contain sensitive information.
	slog.Debug("Successfully extracted authorization header from gRPC request")

	// Parse and validate the JWT token extracted from the authorization header.
	parsedToken, err := oidc.jwtParser.Parse(authHeader, func(token *jwt.Token) (interface{}, error) {
		slog.Info("starting JWT parsing and validation.")

		// Retrieve the key ID from the JWT header and find the corresponding key in the JWKS.
		if keyID, ok := token.Header["kid"].(string); ok {
			return oidc.getKeyWithRetry(requestContext, keyID)
		}
		slog.Error("jwt does not contain a key ID")
		// If the JWT does not contain a key ID, return an error.
		return nil, errors.New("kid must be specified in the token header")
	})
	if err != nil {
		// Log that the token parsing or validation failed
		slog.Error("token parsing or validation failed", "error", err)
		// If token parsing or validation fails, return an error indicating the token is invalid.
		return errors.New(base.ErrorCode_ERROR_CODE_INVALID_BEARER_TOKEN.String())
	}

	// Ensure the token is valid.
	if !parsedToken.Valid {
		// Log that the parsed token was not valid
		slog.Warn("parsed token is invalid")
		// Return an error indicating the invalid token
		return errors.New(base.ErrorCode_ERROR_CODE_INVALID_BEARER_TOKEN.String())
	}

	// Extract the claims from the token.
	claims, ok := parsedToken.Claims.(jwt.MapClaims)
	if !ok {
		// Log that the claims were in an incorrect format
		slog.Warn("token claims are in an incorrect format")
		// Return an error
		return errors.New(base.ErrorCode_ERROR_CODE_INVALID_CLAIMS.String())
	}

	slog.Debug("extracted token claims", "claims", claims)

	// Verify the issuer of the token matches the expected issuer.
	if ok := claims.VerifyIssuer(oidc.IssuerURL, true); !ok {
		// Log that the issuer did not match the expected issuer
		slog.Warn("token issuer is invalid", "expected", oidc.IssuerURL, "actual", claims["iss"])
		// Return an error
		return errors.New(base.ErrorCode_ERROR_CODE_INVALID_ISSUER.String())
	}

	// Verify the audience of the token matches the expected audience.
	if ok := claims.VerifyAudience(oidc.Audience, true); !ok {
		// Log that the audience did not match the expected audience
		slog.Warn("token audience is invalid", "expected", oidc.Audience, "actual", claims["aud"])
		// Return an error
		return errors.New(base.ErrorCode_ERROR_CODE_INVALID_AUDIENCE.String())
	}

	// Log that the token's issuer and audience were successfully validated
	slog.Info("token validation succeeded")

	// If all validations pass, return nil indicating the token is valid.
	return nil
}

// getKeyWithRetry attempts to retrieve the key for the given keyID with retries using a custom backoff strategy.
func (oidc *Authn) getKeyWithRetry(ctx context.Context, keyID string) (interface{}, error) {
	var rawKey interface{}
	var err error

	oidc.mu.Lock()
	now := time.Now()

	// Reset global state if the interval has passed
	if oidc.globalFirstSeen.IsZero() || time.Since(oidc.globalFirstSeen) >= oidc.backoffInterval {
		slog.Info("resetting state as interval has passed or first seen is zero", "keyID", keyID)
		oidc.globalFirstSeen = now
		oidc.globalRetryCount = 0
		oidc.globalRetryKeyIds = make(map[string]bool)
	} else if oidc.globalRetryCount >= oidc.backoffMaxRetries {
		// If max retries reached within the interval, unlock and check keyID once
		slog.Warn("max retries reached within interval, will check keyID once", "keyID", keyID)
		oidc.mu.Unlock()

		// Try to fetch the keyID once
		rawKey, err = oidc.fetchKey(ctx, keyID)
		if err == nil {
			oidc.mu.Lock()
			if _, ok := oidc.globalRetryKeyIds[keyID]; ok {
				// Reset global backoff state if a valid key is found and that key had been retried.
				// Use case would be someone trying to exploit with bad KeyIDs, and along comes a valid KeyID
				// The valid KeyID should not reset the counters for a bad key
				slog.Info("valid key found during backoff period, resetting state", "keyID", keyID)
				oidc.globalRetryCount = 0
				oidc.globalFirstSeen = time.Time{}
				oidc.globalRetryKeyIds = make(map[string]bool)
			}
			oidc.mu.Unlock()
			return rawKey, nil
		}

		// Log the failure and return an error if keyID is not found
		slog.Error("failed to fetch key during backoff period", "keyID", keyID, "error", err)
		return nil, errors.New("too many attempts, backoff in effect")
	}
	oidc.mu.Unlock()

	// Retry mechanism
	retries := 0
	for retries <= oidc.backoffMaxRetries {
		rawKey, err = oidc.fetchKey(ctx, keyID)
		if err == nil {
			if retries != 0 {
				oidc.mu.Lock()
				oidc.globalRetryCount = 0
				oidc.globalFirstSeen = time.Time{}
				oidc.globalRetryKeyIds = make(map[string]bool)
				oidc.mu.Unlock()
			}
			return rawKey, nil
		}
		oidc.mu.Lock()
		initialGlobalRetryCount := oidc.globalRetryCount
		oidc.globalRetryKeyIds[keyID] = true
		if oidc.globalRetryCount > oidc.backoffMaxRetries {
			slog.Error("key ID not found in JWKS due to global retries", "keyID", keyID, "globalRetryCount", oidc.globalRetryCount)
			oidc.mu.Unlock()
			return nil, errors.New("too many attempts, backoff in effect due to global retry count")
		}
		oidc.mu.Unlock()
		if retries > 0 {
			select {
			case <-time.After(oidc.backoffFrequency):
				// Log the wait before retrying
				slog.Info("waiting before retrying", "keyID", keyID, "retries", retries)
			case <-ctx.Done():
				slog.Error("context cancelled during retry", "keyID", keyID)
				return nil, ctx.Err()
			}
		}

		oidc.mu.Lock()
		if oidc.globalRetryCount > initialGlobalRetryCount {
			// Concurrent requests in retry loop at same time, another concurrent request already refreshed the JWKS
			retries++
			slog.Warn("another concurrent request already refreshed the JWKS")
			oidc.mu.Unlock()
			continue
		}

		oidc.globalRetryCount++
		slog.Warn("retrying to fetch JWKS due to error", "keyID", keyID, "retries", retries, "error", err)
		retries++

		if _, refreshErr := oidc.jwksSet.Refresh(ctx, oidc.JwksURI); refreshErr != nil {
			oidc.mu.Unlock()
			slog.Error("failed to refresh JWKS", "error", refreshErr)
			return nil, refreshErr
		}
		// Unlock needs to follow Refresh to ensure that concurrent requests don't make duplicate calls to Refresh
		oidc.mu.Unlock()
	}

	// Mark the global state to prevent further retries for the backoff interval
	oidc.mu.Lock()
	if time.Since(oidc.globalFirstSeen) < oidc.backoffInterval {
		slog.Warn("marking state to prevent further retries", "keyID", keyID)
		oidc.globalRetryCount = oidc.backoffMaxRetries
	}
	oidc.mu.Unlock()

	slog.Error("key ID not found in JWKS after retries", "keyID", keyID)
	return nil, errors.New("key ID not found in JWKS after retries")
}

// fetchKey attempts to fetch the JWKS and retrieve the key for the given keyID.
func (oidc *Authn) fetchKey(ctx context.Context, keyID string) (interface{}, error) {
	// Log the attempt to find the key.
	slog.DebugContext(ctx, "attempting to find key in JWKS", "kid", keyID)

	// Fetch the JWKS from the configured URI.
	jwks, err := oidc.jwksSet.Fetch(ctx, oidc.JwksURI)
	if err != nil {
		// Log an error and return if fetching fails.
		slog.Error("failed to fetch JWKS", "uri", oidc.JwksURI, "error", err)
		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)

unrecognized printf verb 'w'

	}

	// Log a successful fetch of the JWKS.
	slog.InfoContext(ctx, "successfully fetched JWKS")

	// Attempt to find the key in the fetched JWKS using the key ID.
	if key, found := jwks.LookupKeyID(keyID); found {
		var k interface{}
		// Convert the key to a usable format.
		if err := key.Raw(&k); err != nil {
			slog.ErrorContext(ctx, "failed to get raw public key", "kid", keyID, "error", err)
			return nil, fmt.Errorf("failed to get raw public key: %w", err)

unrecognized printf verb 'w'

		}
		// Log a successful retrieval of the raw public key.
		slog.DebugContext(ctx, "successfully obtained raw public key", "key", k)
		return k, nil // Return the public key for JWT signature verification.
	}
	// Log an error if the key ID is not found in the JWKS.
	slog.ErrorContext(ctx, "key ID not found in JWKS", "kid", keyID)
	return nil, fmt.Errorf("kid %s not found", keyID)
}

// Config holds OpenID Connect (OIDC) configuration details.
type Config struct {
	// Issuer is the OIDC provider's unique identifier URL.
	Issuer string `json:"issuer"`
	// JWKsURI is the URL to the JSON Web Key Set (JWKS) provided by the OIDC issuer.
	JWKsURI string `json:"jwks_uri"`
}

// fetchOIDCConfiguration sends an HTTP request to the given URL to fetch the OpenID Connect (OIDC) configuration.
// It requires an HTTP client and the URL from which to fetch the configuration.
func fetchOIDCConfiguration(client *http.Client, url string) (*Config, error) {
	// Send an HTTP GET request to the provided URL to fetch the OIDC configuration.
	// This typically points to the well-known configuration endpoint of the OIDC provider.
	body, err := doHTTPRequest(client, url)
	if err != nil {
		// If there is an error in fetching the configuration (network error, bad response, etc.), return nil and the error.
		return nil, err
	}

	// Parse the JSON response body into an OIDC Config struct.
	// This involves unmarshalling the JSON into a struct that matches the expected fields of the OIDC configuration.
	oidcConfig, err := parseOIDCConfiguration(body)
	if err != nil {
		return nil, err
	}

	// Return the parsed OIDC configuration and nil as the error (indicating success).
	return oidcConfig, nil
}

// doHTTPRequest makes an HTTP GET request to the specified URL and returns the response body.
func doHTTPRequest(client *http.Client, url string) ([]byte, error) {
	// Log the attempt to create a new HTTP GET request
	slog.Debug("creating new HTTP GET request", "url", url)

	// Create a new HTTP GET request.
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		slog.Error("failed to create HTTP request", "url", url, "error", err)
		return nil, fmt.Errorf("failed to create HTTP request for OIDC configuration: %s", err)
	}

	// Log the execution of the HTTP request
	slog.Debug("executing HTTP request", "url", url)

	// Send the request using the configured HTTP client.
	res, err := client.Do(req)
	if err != nil {
		// Log the error if executing the HTTP request fails
		slog.Error("failed to execute HTTP request", "url", url, "error", err)
		return nil, fmt.Errorf("failed to execute HTTP request for OIDC configuration: %s", err)
	}

	// Log the HTTP status code of the response
	slog.Debug("received HTTP response", "status_code", res.StatusCode, "url", url)

	// Ensure the response body is closed after reading.
	defer res.Body.Close()

	// Check if the HTTP status code indicates success.
	if res.StatusCode != http.StatusOK {
		slog.Warn("received unexpected status code", "status_code", res.StatusCode, "url", url)
		return nil, fmt.Errorf("received unexpected status code (%d) while fetching OIDC configuration", res.StatusCode)
	}

	// Log the attempt to read the response body
	slog.Debug("reading response body", "url", url)

	// Read the response body.
	body, err := io.ReadAll(res.Body)
	if err != nil {
		slog.Error("failed to read response body", "url", url, "error", err)
		return nil, fmt.Errorf("failed to read response body from OIDC configuration request: %s", err)
	}

	// Log the successful retrieval of the response body
	slog.Debug("successfully read response body", "url", url, "response_length", len(body))

	// Return the response body.
	return body, nil
}

// parseOIDCConfiguration decodes the OIDC configuration from the given JSON body.
func parseOIDCConfiguration(body []byte) (*Config, error) {
	var oidcConfig Config
	// Attempt to unmarshal the JSON body into the oidcConfig struct.
	if err := json.Unmarshal(body, &oidcConfig); err != nil {
		slog.Error("failed to unmarshal OIDC configuration", "error", err)
		return nil, fmt.Errorf("failed to decode OIDC configuration: %s", err)
	}
	// Log the successful decoding of OIDC configuration
	slog.Debug("successfully decoded OIDC configuration")

	if oidcConfig.Issuer == "" {
		slog.Warn("missing issuer value in OIDC configuration")
		return nil, errors.New("issuer value is required but missing in OIDC configuration")
	}

	if oidcConfig.JWKsURI == "" {
		slog.Warn("missing JWKsURI value in OIDC configuration")
		return nil, errors.New("JWKsURI value is required but missing in OIDC configuration")
	}

	// Log the successful parsing of the OIDC configuration
	slog.Info("successfully parsed OIDC configuration", "issuer", oidcConfig.Issuer, "jwks_uri", oidcConfig.JWKsURI)

	// Return the successfully parsed configuration.
	return &oidcConfig, nil
}