Permify / permify

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"

)

// Authenticator - Interface for oidc authenticator

type Authenticator interface {

	Authenticate(ctx context.Context) error

}

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

	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.

	}

	// 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,

		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 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.

// getKeyWithRetry attempts to retrieve the key for the given keyID with retries using a custom backoff strategy.

func (oidc *Authn) getKeyWithRetry(keyID string, ctx context.Context) (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

	} 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(keyID, ctx)

		if err == nil {

			// Reset global backoff state if a valid key is found

			slog.Info("valid key found during backoff period, resetting state", "keyID", keyID)

			oidc.mu.Lock()

			oidc.globalRetryCount = 0

			oidc.globalFirstSeen = time.Time{}

			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 {

		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()

			}

		}

		rawKey, err = oidc.fetchKey(keyID, ctx)

		if err == nil {

			// Log the successful key fetch and reset global state

			slog.Info("successfully fetched key", "keyID", keyID)

			oidc.mu.Lock()

			oidc.globalRetryCount = 0

			oidc.globalFirstSeen = time.Time{}

			oidc.mu.Unlock()

			return rawKey, nil

		}

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

		retries++

		oidc.mu.Lock()

		oidc.globalRetryCount++

		oidc.mu.Unlock()

		if _, refreshErr := oidc.jwksSet.Refresh(ctx, oidc.JwksURI); refreshErr != nil {

			slog.Error("failed to refresh JWKS", "error", refreshErr)

			return nil, refreshErr

		}

	}

	// 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(keyID string, ctx context.Context) (interface{}, error) {

	// Log the attempt to find the key.

	slog.Debug("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)

	}

	// Log a successful fetch of the JWKS.

	slog.Info("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.Error("failed to get raw public key", "kid", keyID, "error", err)

		}

		// Log a successful retrieval of the raw public key.

		slog.Debug("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.Error("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

}